NVIDIA FastGen: Fast Generation from Diffusion Models

TL;DR

FastGen is an Apache-2.0 PyTorch library from NVIDIA’s GenAIR group that unifies state-of-the-art diffusion distillation methods — trajectory-based (CM, sCM, TCM, MeanFlow) and distribution-based (DMD2, f-Distill, LADD, CausVid, Self-Forcing) — plus fine-tuning (SFT, CausalSFT) and knowledge distillation (KD, CausalKD) recipes under a single configuration interface. The pitch: same network, same data, swap distillation method via a config file; FastGen handles the training loop, FSDP2/AMP/context-parallelism scaling, and reproducibility, and reproduces published one-step FIDs on CIFAR-10 and ImageNet-64. The release ships pre-wired support for EDM / DiT / SD-1.5 / SDXL / Flux / Qwen-Image on the image side and WAN (T2V/I2V/VACE) / CogVideoX / Cosmos Predict2 on the video side, and reports a 14B Wan2.1-T2V model distilled with DMD2 to a few-step generator in 16 H100-hours × 64 GPUs.

Key claims

A unified configuration interface decouples distillation method from network architecture and dataset, so the same backbone (e.g. EDM, DiT, SDXL, WAN) can be distilled with any of the supported recipes by swapping one config file [README §Documentation, §Supported Methods].
Supported one-step / few-step training methods cover both major families: trajectory-based — CM, sCM, TCM, MeanFlow — and distribution-matching — DMD2, f-Distill, LADD, CausVid, Self-Forcing [README §Supported Methods].
Image network coverage at release: EDM, EDM2, DiT, SD 1.5, SDXL, Flux, Qwen-Image; video network coverage: WAN (T2V / I2V / VACE), CogVideoX, Cosmos Predict2 [README §Supported Networks and Data].
FastGen reproduces published one-step FID numbers on CIFAR-10 and ImageNet-64 in a unified codebase for ECT, TCM, sCT, sCD, MeanFlow (trajectory-based) and LADD, DMD2, f-distill (distribution-based) [companion blog §What FastGen offers, Table 1].
A 14B Wan2.1-T2V model is successfully distilled with DMD2 to a two-step generator on 64 H100s in 16 hours, with the student matching the 50-step teacher visually while sampling ~50× faster [companion blog §Scalable and efficient infrastructure, Figure 3].
The library’s training infrastructure supports FSDP2, AMP, context parallelism, Flex attention, KV-cache management, and adaptive finite-difference JVP estimation as built-in primitives for scaling distillation to large models [companion blog §Scalable and efficient infrastructure].
FastGen extends distillation beyond vision: a CorrDiff weather downscaling model in NVIDIA PhysicsNeMo is distilled with ECT to a one-step model achieving 23× faster inference while matching teacher skill/spread, as a “AI for science” datapoint [companion blog §Beyond vision tasks].
For interactive world models, FastGen ships causal-distillation recipes (CausVid, Self-Forcing) plus two novel ingredients: (a) warm-starting causal distillation with trajectory-based methods before applying distribution-based objectives, and (b) Causal SFT via diffusion forcing, training a many-step block-wise autoregressive model first to serve as a teacher for trajectory-based distillation [companion blog §FastGen for interactive world modeling].

Method

FastGen is a PyTorch library, not a paper, so “method” here is the architecture of the library itself. The repo factors into four pluggable axes: methods (the training recipe — sCM / MeanFlow / DMD2 / Self-Forcing / SFT / KD / CausalSFT / CausalKD), networks (the backbone — EDM, DiT, SDXL, Flux, Qwen-Image, WAN, CogVideoX, Cosmos), datasets (CIFAR-10, ImageNet, custom WebDataset loaders), and callbacks (EMA, profiling, gradient clipping, logging). A central Hydra-style configuration system mixes-and-matches these, and a single train.py entry point dispatches to either DDP or FSDP2 multi-GPU training. Trajectory-based losses use adaptive finite-difference JVP estimation as an alternative to native PyTorch JVPs (the same training bottleneck SoFlow circumvented by reformulating the consistency target — see SoFlow: Solution Flow Models for One-Step Generative Modeling).

The interactive-world-modeling track adds two recipes specific to causal distillation: warm-starting (initialise the student with a trajectory-based fit before flipping to distribution-matching), and Causal SFT via diffusion forcing (build an intermediate many-step block-wise AR teacher and then trajectory-distill against it). These address the bidirectional-teacher / autoregressive-student mismatch that limits naive causal distillation.

Results

The companion NVIDIA Developer blog presents Table 1 — one-step FID on CIFAR-10 and ImageNet-64 — comparing FastGen’s unified reimplementations to the originally-published numbers in parentheses, across ECT, TCM, sCT, sCD, MeanFlow, LADD, DMD2, and f-distill. The numbers are presented as matched or close-to-matched, framed as a “transparent benchmarking” contribution.

Headline scaling result: Wan2.1-T2V-14B distilled with DMD2 reaches two-step generation in 16 H100-hours × 64 GPUs, with the student “closely matching” the 50-step teacher’s quality (visual evidence, Figure 3). CorrDiff weather downscaling distilled with ECT to one step, 23× faster than the multi-step teacher, with skill/spread preserved (Figure 2). The README also notes that distilled student checkpoints for CIFAR-10 and ImageNet are planned but not yet shipped.

No standalone “FastGen invented X” claims — this is explicitly a library that re-implements existing methods inside one framework, not a new algorithm. The reproducibility and the 14B/64-GPU distillation demonstration are the artifact’s main contributions.

Why it’s interesting

Three angles for Luma. (1) FastGen is the closest thing to a reference implementation lookup table for the diffusion-distillation landscape — every method on the team’s reading list (CM, sCM, MeanFlow, DMD2, f-Distill, LADD, CausVid, Self-Forcing) plus the fine-tuning and KD baselines lives in one Apache-2.0 repo with comparable hyperparameters, making it the natural starting point for “do A and B compose?” experiments without having to port across N codebases. (2) It bakes FSDP2 + context parallelism + Flex attention + adaptive-FD JVPs into the default training loop, so the scaling story (14B Wan2.1 distilled on 64 H100s) is reproducible rather than aspirational. (3) The causal- distillation track is directly aligned with the autoregressive video work the team has been filing — both warm-starting and Causal SFT via diffusion forcing are training-time complements to the Context Forcing / PFP / LongVie 2 lineage of long-context AR video. The 16-hour / 64-GPU figure for distilling a 14B T2V model is also a useful anchor for “what does it cost to distill a Luma-sized video DiT in 2026”.