Nvidia y Hugging Face se unen para hacer aburrido el entrenamiento distribuido de difusión

Los modelos de difusión siguen creciendo y entrenarlos se vuelve cada vez más incómodo. FLUX.1-dev pesa 12 mil millones de parámetros. La variante de 14 mil millones de parámetros de Wan 2.1 necesita paralelismo tensorial solo para caber en un H100. HunyuanVideo 1.5 funciona con 13 mil millones. Escalar la inferencia para estos modelos es relativamente bien entendido. Escalar el entrenamiento no lo es, especialmente si quieres pasar de un experimento LoRA con una sola GPU a un ajuste fino completo en un clúster sin reescribir tu script de entrenamiento.

Esa brecha es lo que Nvidia y Hugging Face están abordando con una nueva integración entre NeMo Automodel y la biblioteca Diffusers, ahora documentada en la guía de entrenamiento de Diffusers y publicada bajo Apache 2.0. La propuesta es directa: apuntas a cualquier modelo compatible con Diffusers en el Hub, y el mismo flujo de trabajo basado en YAML maneja el fragmentado FSDP2, el paralelismo tensorial, el paralelismo de contexto, el almacenamiento en caché latente y el bucket multirresolución. Sin paso de conversión. Sin formato de entrenamiento separado. El punto de control sale en formato Diffusers y se carga directamente en un pipeline de difusión para inferencia.

Para los equipos que han estado creando scripts de entrenamiento personalizados alrededor de cada nuevo lanzamiento de modelo, ese empaquetado importa más que cualquier número de rendimiento individual.

¿Qué cambia realmente?

Las ganancias prácticas se dividen en tres partes. Primero, sin conversión de punto de control. Un punto de control de NeMo ajustado es un punto de control de Diffusers, la misma estructura de carpetas, las mismas clases de modelo. Las herramientas posteriores (cuantización, compilación, samplers personalizados, fusión LoRA) siguen funcionando sin un paso adaptador. Segundo, soporte rápido para nuevos modelos. Cuando una nueva arquitectura de difusión llega a Diffusers, habilitarla en NeMo Automodel requiere un manejador de preprocesamiento de datos y un adaptador de modelo en lugar de un script de entrenamiento personalizado completo. El resto del stack de recetas, paralelismo, bucket, checkpointing, generación, se mantiene sin cambios.

Tercero, tanto el ajuste fino completo como el estilo LoRA PEFT son manejados por la misma estructura de recetas. La configuración YAML selecciona cuál ejecutar; la ruta del código es compartida. Esto significa que un profesional puede comenzar con LoRA en un solo nodo para validar la calidad de los datos, luego escalar a un ajuste fino completo en ocho GPU (o ochenta) cambiando algunos campos de configuración y una declaración de paralelismo, no intercambiando scripts de entrenamiento.

La capa de referencia

Los benchmarks publicados en 8x H100 80GB dan una idea de dónde se sitúa la biblioteca. El ajuste fino completo de FLUX.1-dev a resolución 512x512 funciona a 35.5 imágenes por segundo con un tamaño de lote global de 32 y asignación máxima de GPU de 63.9 GiB. LoRA con rango 64 y DDP lo lleva a 53.7 imágenes por segundo, consumiendo 67.4 GiB. Para texto a video, el ajuste fino completo de Wan 2.1 14B en clips de 49 fotogramas a 512x512 funciona a aproximadamente 2.1 clips por segundo con checkpointing de activación activado, alcanzando un máximo de 33.4 GiB por GPU. La variante 1.3B cabe fácilmente, 6.1 GiB máximo, sin necesidad de checkpointing de activación.

Esos números son promedios de estado estacionario con escrituras de puntos de control desactivadas y lotes completos forzados. Las ejecuciones del mundo real con checkpointing periódico serán ligeramente más lentas, pero la brecha entre LoRA de una sola GPU y el ajuste fino completo distribuido es lo suficientemente estrecha como para que el principal cuello de botella para la mayoría de los equipos sea la preparación de datos y la evaluación, no el rendimiento del entrenamiento.

Un recorrido concreto

La documentación incluye un ajuste fino completo de extremo a extremo de FLUX.1-dev en el conjunto de datos de tarot Rider-Waite de 78 imágenes. El flujo de trabajo es: pre-codificar el conjunto de datos en latentes VAE y embeddings de texto en caché, luego lanzar el entrenamiento usando el YAML FLUX existente con anulaciones de línea de comandos para las rutas del conjunto de datos y la configuración de ejecución. Después de 200 pasos de optimización, el punto de control produce salidas que combinan el contenido del prompt de astronauta con las paletas vintage y los contornos de tinta aprendidos de las cartas del tarot, una adaptación de dominio controlada en lugar de una toma de control del modelo base.

El token de activación trtcrd funciona como un interruptor de estilo. Con él, las generaciones cambian a campos de color plano crema, rojo y negro. Sin él, la misma semilla y prompt producen un resultado fotográfico. Ese tipo de control fino es exactamente lo que debería ser el ajuste fino de dominio, y el hecho de que tomó un solo archivo de configuración y ninguna modificación del código fuente del modelo es el punto.

La pieza faltante

NeMo Automodel actualmente solo admite modelos de coincidencia de flujo. Eso deja fuera los modelos de difusión de tiempo discreto que todavía se usan ampliamente en la comunidad, arquitecturas de clase Stable Diffusion 3 y SDXL, por ejemplo. La biblioteca también incluye orquestación multi-nodo SLURM pero aún no soporte para Kubernetes, lo que limita su atractivo para equipos que trabajan en plataformas en la nube gestionadas.

La API de recetas Pythonic planeada, que permitiría a los usuarios componer los mismos componentes desde Python con tipos en lugar de YAML, todavía está listada como un lanzamiento futuro. Para los equipos que necesitan integrar el entrenamiento en marcos de gestión de experimentos existentes, esa ruta programática importará más que el inicio rápido YAML.

Nada de eso socava lo que realmente está aquí. La integración Diffusers-NeMo resuelve un problema específico, el entrenamiento distribuido de difusión de nivel de producción sin el andamiaje habitual, y lo resuelve limpiamente. Los números de rendimiento son sólidos, el flujo de trabajo es reproducible y el artefacto funciona con las herramientas existentes. Para los laboratorios que han estado ejecutando scripts ad hoc en clústeres, este es probablemente el camino de menor resistencia hacia una configuración que escala.