FastGen: NVIDIA’s Open‑Source Framework Accelerating Diffusion Models

What is FastGen?

FastGen is a PyTorch‑based framework released by NVIDIA that lets researchers and developers build fast generative models from diffusion architectures. By combining a collection of distillation methods, efficient training recipes, and a modular configuration system, FastGen can turn a 12‑Billion‑parameter diffusion model into a lightweight, real‑time‑capable student network.

Key attributes:

• Speed‑of‑execution – Turn hours of training into minutes and billions of parameters into 10‑B students.
• Versatility – Supports T2I, I2V, V2V, and other modalities across a variety of backbone models (EDM, SDXL, Flux, CogVideoX, etc.).
• Extensibility – Plug‑in new datasets, networks, and distillation algorithms with minimal friction.
• Multi‑GPU & FSDP2 – Built‑in support for DDP and FairScale‑based model sharding.

Why FastGen is a Must‑Have Tool

Diffusion models have become the de‑facto standard for image and video generation due to their high fidelity and controllability. However, their training cost can be prohibitive. FastGen addresses this by:

1. Distillation techniques – Consistency models, distribution‑matching (DMD2, LADD), self‑forcing, and knowledge distillation compress large teachers into efficient students.
2. Dynamic batching – Automatic gradient accumulation to meet global batch size even on small GPUs.
3. Configuration hierarchy – Hydra‑style configs that separate experiment details from method‑specific hyper‑parameters.
4. Reproducibility – Every run outputs a resolved config, checkpoints, and a W&B run ID.
5. Rich documentation – Every component has a dedicated README.

Quick Start Guide

Below is the minimal workflow to get a FastGen model up and running:

1. Clone the repo

   git clone https://github.com/NVlabs/FastGen.git
   cd FastGen
   

2. Create a conda environment

   conda create -y -n fastgen python=3.12.3
   conda activate fastgen
   

3. Install in editable mode

   pip install -e .
   

4. Download dataset and reference models (CIFAR‑10 example)

   python scripts/download_data.py --dataset cifar10
   

5. Run a training experiment

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

6. Inspect the results – W&B console will pop up; logs and checkpoints are stored under FASTGEN_OUTPUT/fastgen/cifar10/....

Multi‑GPU Training

FastGen supports torchrun for DDP and FairScale’s FSDP2 for model sharding.

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

For FSDP2 replace trainer.ddp=True with trainer.fsdp=True.

Inference

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

Results can be evaluated with FID or other metrics via the scripts/README.md guidelines.

How to Contribute

FastGen welcomes all contributors. Suggested entry points:

• Add new distillation methods – Implementation templates live in fastgen/methods/.
• Integrate new datasets – Follow the pattern in fastgen/datasets/.
• Improve docs – Anything that helps users run or extend FastGen.
• Report bugs – Check the Issues tab for open tasks.

The main development flow: fork → create a feature branch → open PR → review.

Future Roadmap

• Pre‑trained student checkpoints – For ImageNet and CIFAR‑10 will be released soon.
• Expanded modality support – Video‑to‑image and image‑to‑audio extensions.
• Hardware‑agnostic optimizations – Better support for Apple Silicon and emerging accelerators.

Final Thoughts

FastGen exemplifies how an open‑source framework can democratise cutting‑edge AI. Its modular design means you can drop into a pre‑built experimental pipeline, experiment with various distillation strategies, or extend the library to accommodate new research directions. Whether you’re a researcher pushing the envelope of diffusion theory or a startup seeking efficient generative models, FastGen offers a proven, battle‑tested base to accelerate your projects.

Quick Tip: If you only need inference, the scripts/inference/ folder has lightweight utilities that can be run on a single GPU with minimal memory.

Feel free to star the GitHub repo, join the discussion, and share your experiments. NVIDIA’s FastGen is open‑source, community‑driven, and ready to power your next generation of generative AI.