> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/fastgen/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/fastgen/_mcp/server.

# Get started

[View source on GitHub](https://github.com/NVlabs/FastGen/blob/7889625364c56fa02d2a54c1683c99a98bc891c3/README.md)

## Setup

**Recommended:** Use the provided Docker container for a consistent environment. See [CONTRIBUTING.md](https://github.com/NVlabs/FastGen/blob/7889625364c56fa02d2a54c1683c99a98bc891c3/CONTRIBUTING.md) for Docker setup instructions. Otherwise, create a new [conda](https://www.anaconda.com/docs/getting-started/miniconda/install) environment with `conda create -y -n fastgen python=3.12.3 pip; conda activate fastgen`.

### Installation

```bash
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:

```bash
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](https://wandb.ai/settings) 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](/fastgen/configuration#environment-variables).

## Quick Start

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

```bash
python scripts/download_data.py --dataset cifar10
```

For other datasets and models, see [fastgen/networks/README.md](/fastgen/networks) and [fastgen/datasets/README.md](/fastgen/datasets).

### Basic Training

```bash
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](https://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](/fastgen/configuration#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:

```bash
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

```bash
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/inference).
FastGen-distilled models can also be served with:

* [FlashDreams](/fastgen/inference#flashdreams-example-self-forcing-wan-21-t2v) for optimized streaming inference of Self-Forcing models.
* [vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/models/dmd2/) for DMD2 model inference.

### Command-Line Overrides

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

```bash
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](/fastgen/configuration#command-line-arguments).