FastGen: Marco de código abierto de NVIDIA que acelera los modelos de difusión

¿Qué es FastGen?

FastGen es un framework basado en PyTorch lanzado por NVIDIA que permite a investigadores y desarrolladores construir modelos generativos rápidos a partir de arquitecturas de difusión. Al combinar una colección de métodos de distilación, recetas de entrenamiento eficientes y un sistema de configuración modular, FastGen puede convertir un modelo de difusión de 12 mil millones de parámetros en una red estudiante ligera y capaz de tiempo real.

Atributos clave:

• Velocidad de ejecución – Convierte horas de entrenamiento en minutos y billones de parámetros en 10 mil millones de estudiantes.
• Versatilidad – Soporta T2I, I2V, V2V y otras modalidades a través de una variedad de modelos base (EDM, SDXL, Flux, CogVideoX, etc.).
• Extensibilidad – Integra nuevos conjuntos de datos, redes y algoritmos de distilación con mínima fricción.
• Multi‑GPU y FSDP2 – Soporte incorporado para DDP y model sharding basado en FairScale.

Por qué FastGen es una herramienta indispensable

Los modelos de difusión se han convertido en el estándar de facto para la generación de imágenes y videos debido a su alta fidelidad y controlabilidad. Sin embargo, su coste de entrenamiento puede ser prohibitivo. FastGen aborda esto mediante:

1. Técnicas de distilación – Modelos de consistencia, equivalencia de distribución (DMD2, LADD), autocontención y distilación de conocimiento comprimen grandes maestros en estudiantes eficientes.
2. Batching dinámico – Acumulación automática de gradientes para alcanzar el tamaño de lote global incluso en GPUs pequeñas.
3. Jerarquía de configuración – Configuraciones estilo Hydra que separan detalles experimentales de hiperparámetros específicos del método.
4. Reproducibilidad – Cada ejecución genera una configuración resuelta, puntos de control y un ID de ejecución W&B.
5. Documentación rica – Cada componente cuenta con un README dedicado.

Guía de inicio rápido

A continuación se muestra el flujo de trabajo mínimo para poner en marcha un modelo FastGen:

1. Clonar el repositorio

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

2. Crear un entorno conda

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

3. Instalar en modo editable

   pip install -e .
   

4. Descargar conjunto de datos y modelos de referencia (ejemplo CIFAR‑10)

   python scripts/download_data.py --dataset cifar10
   

5. Ejecutar un experimento de entrenamiento

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

6. Inspeccionar los resultados – La consola de W&B aparecerá; los registros y puntos de control se almacenan bajo FASTGEN_OUTPUT/fastgen/cifar10/....

Entrenamiento multi‑GPU

FastGen soporta torchrun para DDP y FSDP2 de FairScale para sharding de modelo.

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

Para FSDP2 sustituye trainer.ddp=True por trainer.fsdp=True.

Inferencia

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

Los resultados pueden evaluarse con FID u otras métricas siguiendo las directrices de scripts/README.md.

Cómo contribuir

FastGen da la bienvenida a todos los colaboradores. Puntos de entrada sugeridos:

• Añadir nuevos métodos de distilación – Plantillas de implementación se encuentran en fastgen/methods/.
• Integrar nuevos conjuntos de datos – Sigue el patrón en fastgen/datasets/.
• Mejorar la documentación – Cualquier cosa que ayude a los usuarios a ejecutar o ampliar FastGen.
• Reportar errores – Consulta la pestaña de Issues para tareas abiertas.

El flujo de desarrollo principal: fork → crear rama de característica → abrir PR → revisión.

Hoja de ruta futura

• Puntos de control de estudiantes preentrenados – Para ImageNet y CIFAR‑10 se lanzarán pronto.
• Soporte ampliado de modalidades – Extensiones de video‑a‑imagen y imagen‑a‑audio.
• Optimizaciones agnósticas al hardware – Mejor soporte para Apple Silicon y aceleradores emergentes.

Reflexiones finales

FastGen ejemplifica cómo un framework de código abierto puede democratizar la IA de vanguardia. Su diseño modular permite que te sumes a una tubería experimental preconstruida, experimentes con distintas estrategias de distilación o amplíes la biblioteca para nuevas líneas de investigación. Sea que estés investigando los límites de la teoría de difusión o dirigiendo una startup que busca modelos generativos eficientes, FastGen ofrece una base probada y battle‑tested para acelerar tus proyectos.

Consejo rápido: si solo necesitas inferencia, la carpeta scripts/inference/ tiene utilidades ligeras que pueden ejecutarse en una sola GPU con memoria mínima.

No dudes en marcar el repositorio de GitHub, unirte a la discusión y compartir tus experimentos. NVIDIA’s FastGen es de código abierto, impulsado por la comunidad y listo para impulsar la próxima generación de IA generativa.