Entrene y ajuste modelos de video e imagen a escala con NeMo Automodel y 🤗 Diffusers

Resumen ejecutivo

La colaboración entre NVIDIA y Hugging Face promete simplificar y escalar el entrenamiento de modelos de difusión (diffusion models) para imagen y video. NeMo Automodel, una librería open source de NVIDIA, se integra de forma nativa con la familia Diffusers de Hugging Face para ofrecer recetas de entrenamiento distribuidas, soporte para fine-tuning tanto completo como con LoRA, y una cadena de trabajo que evita conversiones de checkpoints.

Para equipos y tomadores de decisión en América Latina, esto significa menores barreras técnicas para adaptar modelos de última generación a casos de uso locales —desde generación de contenido multimedia hasta investigación de medios sintéticos— aprovechando arquitecturas y herramientas ya consolidadas en la comunidad.

¿Qué es NeMo Automodel y por qué importa?

NeMo Automodel es una librería de entrenamiento open source basada en PyTorch y DTensor, parte del ecosistema NVIDIA NeMo. Está diseñada con dos principios clave pensados para el ecosistema Diffusers:

• Integración nativa con Hugging Face: basta apuntar pretrained_model_name_or_path a cualquier modelo en formato Diffusers alojado en el Hub para iniciar entrenamiento. NeMo Automodel usa las clases y pipelines de Diffusers para cargar modelos y para generación, y los checkpoints resultantes son compatibles directamente con DiffusionPipeline.

• Escalabilidad sin reescrituras de código: las estrategias de paralelismo (FSDP2, tensor parallel, expert parallel, context parallel, pipeline parallel) se activan mediante configuraciones, no modificando el código del modelo. Esto permite llevar el mismo flujo de trabajo desde un nodo con una GPU hasta clusters de cientos.

En la práctica, NeMo Automodel aprovecha flow matching como objetivo de entrenamiento, realiza entrenamiento en espacio latente (usando latentes VAE pre-encodificados) y aplica bucketed dataloading multiresolución para mejorar throughput.

Modelos soportados (recetas incluidas)

NeMo Automodel ya incluye recetas listas para fine-tuning de varios modelos abiertos que conviene conocer:

• Wan 2.1 (text-to-video): Wan-AI/Wan2.1-T2V-1.3B-Diffusers y Wan-AI/Wan2.1-T2V-14B-Diffusers. La versión de 1.3B cabe en una sola A100 de 40 GB; la versión de 14B tiene receta LoRA.
• Wan 2.2 (text-to-video) — Wan-AI/Wan2.2-T2V-A14B-Diffusers: arquitectura MoE con 27B total y 14B activos por paso.
• FLUX.1-dev (text-to-image) — black-forest-labs/FLUX.1-dev: modelo de 12B con receta disponible.
• FLUX.2-dev (text-to-image) — black-forest-labs/FLUX.2-dev: 32B.
• HunyuanVideo 1.5 (text-to-video) — hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v: 13B con receta.
• Qwen-Image (text-to-image) — Qwen/Qwen-Image: 20B (MMDiT) con receta.

Estas recetas aparecen en examples/diffusion/finetune del repositorio y permiten tanto fine-tuning completo como técnicas de ajuste de parámetros (PEFT) tipo LoRA cuando están disponibles.

Ventajas clave de la integración

1. Sin conversiones de checkpoints: los pesos preentrenados del Hub funcionan tal cual. Los checkpoints finetuneados se cargan directamente en DiffusionPipeline para inferencia o se suben al Hub sin conversiones intermedias.

2. Soporte rápido a nuevos modelos: cuando aparece un nuevo modelo en Diffusers, integrarlo en NeMo Automodel requiere una pequeña adaptación (handler de preprocesamiento y un adapter de modelo), sin rehacer todo el stack de entrenamiento.

3. Opciones de fine-tuning flexibles: la misma infraestructura maneja fine-tuning completo y LoRA, dejando elegir entre máxima calidad o eficiencia de recursos.

4. Escalabilidad avanzada: sharding como FSDP2, paralelismo tensorial, context y pipeline, orquestación multi-nodo (SLURM hoy; Kubernetes en camino), y bucketed dataloading multiresolución, habilitan entrenar modelos grandes como FLUX.1-dev o HunyuanVideo.

Flujo de trabajo típico para fine-tuning

El flujo recomendado reutiliza las recetas y configuraciones YAML incluidas, aplicando ajustes por línea de comandos cuando sea necesario. Resumen de pasos:

1. Preparar el entorno

• NVIDIA provee un contenedor Docker listo: nvcr.io/nvidia/nemo-automodel:26.06 con PyTorch, TransformerEngine y dependencias CUDA compiladas.
• Alternativamente, se puede instalar con pip: pip3 install nemo-automodel o desde el repositorio fuente.

2. Pre-encode del dataset

La receta de entrenamiento consume latentes VAE cacheados y embeddings de texto en lugar de volver a codificar imágenes cada paso. Esto reduce significativamente el cómputo y la memoria durante entrenamiento. En el ejemplo del proyecto se usa el set de 78 imágenes Rider–Waite (tarot) pre-procesadas como demostración.

3. Lanzar el entrenamiento

Las configuraciones YAML provistas (por ejemplo para FLUX.1-dev) contienen los parámetros de paralelismo y data loading; los equipos solo necesitan ajustar rutas de datos y hiperparámetros específicos para su tarea. Cambiar entre FSDP2, tensor parallel o pipeline parallel es una cuestión de declarar la configuración apropiada.

4. Generación y despliegue

Los checkpoints finales son inmediatamente utilizables con Diffusers para generación, o pueden subirse al Hub. Herramientas downstream como cuantización, compilación, adaptadores LoRA y samplers personalizados siguen siendo compatibles.

Consideraciones prácticas para equipos en América Latina

• Acceso a hardware: muchos grupos en la región usan proveedores cloud o clusters institucionales con GPUs NVIDIA (A100, H100). La posibilidad de escalar desde una sola GPU hasta clusters grandes facilita iterar en prototipos y luego pasar a producción.

• Costos y estrategia: usar LoRA para fases de experimentación y validación puede reducir costos antes de comprometerse con fine-tuning completo en un cluster.

• Integración en pipelines empresariales: la compatibilidad con formatos Diffusers y la ausencia de conversiones de checkpoints simplifican la adopción por equipos que ya usan herramientas del ecosistema Hugging Face.

Cómo empezar hoy

• Probar el contenedor Docker nvcr.io/nvidia/nemo-automodel:26.06 o instalar nemo-automodel por pip.
• Revisar las recetas en examples/diffusion/finetune del repositorio de NeMo Automodel.
• Empezar con modelos que quepan en su infraestructura (por ejemplo Wan 2.1 1.3B para una sola A100 de 40 GB) y escalar conforme sea necesario.

Conclusión

La integración de NeMo Automodel con Hugging Face Diffusers reduce fricciones técnicas y ofrece un camino escalable y reproducible para entrenar y adaptar modelos de difusión de imagen y video. Para organizaciones y equipos latinoamericanos interesados en IA generativa, representa una opción sólida para llevar modelos de investigación a entornos productivos, con soporte para fine-tuning eficiente, paralelismo avanzado y compatibilidad directa con el ecosistema Diffusers.

Recursos

• Repositorio y guías de NeMo Automodel (instalación y ejemplos en el repositorio oficial).
• Documentación de entrenamiento de Diffusers para detalles de la integración y workflows.