Get started

View source on GitHub

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.