FastGen：NVIDIA 的开源框架，加速扩散模型

FastGen 是什么？

NVIDIA 公开的基于 PyTorch 的框架 FastGen 使研究者和开发者能够从扩散架构构建高速生成模型。通过结合多种蒸馏方法、有效的训练方案和模块化配置系统，FastGen 可以将 120 亿参数的扩散模型转变为轻量化、实时可用的学生网络。

主要属性：

• 执行速度 – 将训练数小时压缩为分钟，将数十亿参数压缩为 10 亿级的学生模型。
• 多功能性 – 支持 T2I、I2V、V2V 等多种模态，涵盖多种主干模型（EDM、SDXL、Flux、CogVideoX 等）。
• 可扩展性 – 以最小的阻力接入新数据集、网络和蒸馏算法。
• 多 GPU 与 FSDP2 – 内置 DDP 与基于 FairScale 的模型分片支持。

为什么 FastGen 是必备工具

扩散模型凭借其高保真度与可控性已成为图像与视频生成的事实标准。然而，其训练成本往往颇高。FastGen 通过以下方式解决这一难题：

1. 蒸馏技术 – 一致性模型、分布匹配（DMD2、LADD）、自强及知识蒸馏将大型教师压缩为高效学生。
2. 动态批处理 – 自动梯度累积，即使在小显存 GPU 上亦可满足全局 batch 大小。
3. 配置层次 – 类 Hydra 的配置方式将实验细节与方法专属超参数分离。
4. 可复现 – 每次运行输出已解析的配置、检查点及 W&B 运行 ID。
5. 丰富文档 – 每一组件都有专门的 README。

快速入门指南

以下是启动 FastGen 模型的最小工作流：

1. 克隆仓库

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

2. 创建 Conda 环境

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

3. 以可编辑模式安装

   pip install -e .
   

4. 下载数据集与参考模型（CIFAR‑10 示例）

   python scripts/download_data.py --dataset cifar10
   

5. 运行训练实验

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

6. 检查结果 – W&B 控制台将弹出；日志与检查点存放于 FASTGEN_OUTPUT/fastgen/cifar10/...。

多 GPU 训练

FastGen 支持 torchrun 用于 DDP 以及 FairScale 的 FSDP2 模型分片。

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

若使用 FSDP2，请将 trainer.ddp=True 替换为 trainer.fsdp=True。

推理

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

结果可按 FID 或其他指标在 scripts/README.md 指导下评估。

如何贡献

FastGen 欢迎所有贡献者。建议的入口点：

• 新增蒸馏方法 – 实现模板位于 fastgen/methods/。
• 集成新数据集 – 参照 fastgen/datasets/ 的模式。
• 改进文档 – 任何能帮助用户运行或扩展 FastGen 的内容。
• 报告漏洞 – 请查看 Issues 页面中的待办任务。

主要开发流程：fork → 创建功能分支 → 提交 PR → 代码审查。

未来路线图

• 预训练学生检查点 – ImageNet 与 CIFAR‑10 的预训练模型即将发布。
• 扩展模态支持 – 视频→图像与图像→音频扩展。
• 硬件无关优化 – 进一步提升 Apple Silicon 与新兴加速器的支持。

最后思考

FastGen 展示了开源框架如何让最前沿的 AI 民主化。其模块化设计使得您可以直接使用预构建的实验流水线，尝试多种蒸馏策略，或扩展库以适配新研究方向。无论您是推动扩散理论边界的研究者，还是寻求高效生成模型的创业公司，FastGen 都提供了经过验证的、经战斗检验的基础，助您加速项目落地。

快速提示：若您只需要推理功能，scripts/inference/ 目录下提供了轻量级工具，可在单 GPU 上以最小内存运行。

欢迎点亮 GitHub 仓库、加入讨论并分享您的实验。NVIDIA 的 FastGen 是开源、社区驱动且已准备好为您的下一代生成式 AI 加速。