AIBit-Descubre proyectos de código abierto AIBit-Descubre proyectos de código abierto

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

Descubre FastGen, el framework de PyTorch de código abierto de NVIDIA diseñado para acelerar el entrenamiento e inferencia de modelos de difusión mediante técnicas avanzadas de distilación y aceleración. Conoce cómo FastGen respalda modelos a gran escala, múltiples modalidades (T2I, I2V, V2V) y una tubería de entrenamiento modular. El artículo cubre instalación, inicio rápido, entrenamiento multi‑GPU, inferencia y cómo los desarrolladores pueden contribuir o ampliar la biblioteca. Sumérgete en el código y descubre por qué FastGen es un cambio de juego para cualquiera que construya IA generativa de alto rendimiento.

¿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.