Escalando el Ajuste Fino de Modelos de Imagen y Video con NVIDIA NeMo Automodel y Hugging Face Diffusers

Los modelos de difusión han impulsado algunos de los lanzamientos open source más emocionantes de los últimos dos años — piensa en FLUX.1-dev para texto a imagen, o Wan 2.1 y HunyuanVideo para texto a video. La biblioteca 🤗 Diffusers se ha convertido en el hogar por defecto de estos modelos, ofreciendo a investigadores y desarrolladores una interfaz única y consistente para inferencia, adaptación y composición de pipelines.

El entrenamiento y ajuste fino de estos modelos también está en auge, y eso requiere herramientas que manejen fragmentación eficiente en memoria, caché de latentes, agrupación multiresolución y configuraciones que escalen suavemente desde una GPU hasta cientos.

Aquí entra la biblioteca open source NVIDIA NeMo Automodel. Hoy destacamos la colaboración entre NVIDIA y Hugging Face que trae entrenamiento distribuido de difusión de grado de producción a cualquier modelo en formato Diffusers en el Hugging Face Hub — sin conversión de checkpoints, sin reescritura de modelos para modelos nuevos. La integración está documentada en la guía de entrenamiento de Diffusers y es completamente open source bajo Apache 2.0.

¿Qué es NeMo Automodel?

NeMo Automodel es una biblioteca de entrenamiento open source nativa de PyTorch DTensor, parte del framework NVIDIA NeMo. Está construida alrededor de dos principios de diseño que importan para el ecosistema Diffusers:

• Nativo de Hugging Face. Apunta pretrained_model_name_or_path a cualquier ID de modelo Diffusers en el Hub y comienza a entrenar. NeMo Automodel usa clases de modelo Diffusers (ej., WanTransformer3DModel) para cargar y pipelines Diffusers (WanPipeline) para generar. Los checkpoints viajan limpios de vuelta al ecosistema Diffusers.

• Un programa, cualquier escala. Las recetas y scripts de entrenamiento se pueden modificar fácilmente para adaptarse a cualquier escala. El paralelismo es una elección de configuración, no una reescritura de código — cambia entre FSDP2, paralelismo tensorial, paralelismo de expertos, paralelismo de contexto y paralelismo de pipeline declarando configuraciones, no reescribiendo modelos.

AutoModel actualmente solo soporta modelos de flow matching. Internamente, usa flow matching como objetivo de entrenamiento, con entrenamiento en espacio latente (a través de salidas VAE pre-codificadas) y carga de datos con agrupación multiresolución para acelerar el rendimiento.

Modelos de difusión soportados

La integración de NeMo Automodel incluye recetas de ajuste fino listas para usar para los modelos de difusión abiertos a continuación. La lista refleja las recetas actualmente en examples/diffusion/finetune.

ModeloID en Hugging FaceTareaParámetrosReceta LoRAWan 2.1 T2V 1.3B / 14BWan-AI/Wan2.1-T2V-1.3B-Diffusers / Wan-AI/Wan2.1-T2V-14B-DiffusersTexto a video1.3B (cabe en una sola A100 de 40GB) / 14BSíWan 2.2 T2V A14BWan-AI/Wan2.2-T2V-A14B-DiffusersTexto a video27B total (MoE), 14B activos por pasoNoFLUX.1-devblack-forest-labs/FLUX.1-devTexto a imagen12BSíFLUX.2-devblack-forest-labs/FLUX.2-devTexto a imagen32BSíHunyuanVideo 1.5hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2vTexto a video13BSíQwen-ImageQwen/Qwen-ImageTexto a imagen20B (MMDiT)Sí

Lo que desbloquea esta colaboración

Para los usuarios de Diffusers, las ganancias prácticas se dividen en algunas capacidades concretas.

Sin conversión de checkpoints. Los pesos preentrenados del Hub funcionan de inmediato. No hay un "formato de entrenamiento" separado al que convertir, y luego convertir de vuelta. Tu checkpoint ajustado fino se carga directamente en un DiffusionPipeline para inferencia, o de vuelta al Hub para compartir. Las herramientas posteriores — cuantización, compilación, adaptadores LoRA, muestreadores personalizados — siguen funcionando.

Camino rápido para soporte de nuevos modelos. Cuando un nuevo modelo de difusión llega a Diffusers, habilitarlo en NeMo Automodel requiere una adición de código pequeña y contenida — un manejador de preprocesamiento de datos y un adaptador de modelo — en lugar de un script de entrenamiento personalizado completo. El resto de la pila de recetas (FSDP2, carga de datos con agrupación, checkpointing, generación) se mantiene sin cambios, y el mismo flujo de trabajo basado en YAML se aplica.

Ajuste fino completo y eficiente en parámetros. Se soportan tanto el ajuste fino completo como el PEFT estilo LoRA, así que puedes elegir entre máxima calidad (FT completo en un clúster grande) o máxima eficiencia (LoRA en un solo nodo). La misma estructura de receta maneja ambos.

Entrenamiento escalable que va más allá de lo que ofrecen los scripts integrados. NeMo Automodel añade esquemas de fragmentación como FSDP2, paralelismo tensorial, de contexto y de pipeline, orquestación multi-nodo (SLURM hoy, Kubernetes próximamente) y agrupación multiresolución. Estas capacidades hacen posible entrenar modelos más grandes como FLUX.1-dev (12B) y HunyuanVideo (13B).

Un vistazo al flujo de trabajo de ajuste fino

En esta sección, recorremos el flujo de trabajo típico para ajustar fino cualquiera de los modelos soportados. La forma recomendada de instalar Automodel es el contenedor Docker de NeMo Automodel (nvcr.io/nvidia/nemo-automodel:26.06), que incluye PyTorch, TransformerEngine y otras dependencias compiladas con CUDA pre-construidas. Alternativamente, instala con pip3 install nemo-automodel o desde fuente (pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git); consulta la guía de instalación para todas las opciones.

Esta guía recorre un ajuste fino de transformador completo de FLUX.1-dev en el conjunto de datos de tarot Rider–Waite de 78 cartas, y luego genera desde el checkpoint resultante. Reutiliza los archivos YAML de configuración ya incluidos y aplica configuraciones específicas de la ejecución como anulaciones de línea de comandos, por lo que no se requieren nuevos archivos de configuración.

1. Pre-codificar el conjunto de datos

La receta de difusión consume latentes VAE en caché y embeddings de texto en lugar de codificar imágenes fuente durante cada paso de entrenamiento. Transmite las 78 imágenes de Rider–Waite directamente desde Hugging Face y distribuye el preprocesamiento entre todas las GPUs visibles:

uv run --locked --no-default-groups \
  --extra diffusion \
  --extra diffusion-media \
python -m tools.diffusion.preprocessing_multiprocess image \
  --dataset_name multimodalart/1920-raider-waite-tarot-public-domain \
  --dataset_media_column image \
  --dataset_caption_column caption \
  --dataset_streaming \
  --max_images 78 \
  --output_dir /cache/flux_tarot \
  --processor flux \
  --model_name black-forest-labs/FLUX.1-dev \
  --max_pixels 245760

Los pies de foto ya contienen el token disparador trtcrd. Con este presupuesto de píxeles y la relación de aspecto vertical del conjunto de datos, el preprocesamiento asigna las muestras al bucket de 384×640 utilizado en la ejecución de demostración.

Para el entrenamiento de imágenes, el preprocesamiento produce archivos de caché .pt y metadatos fragmentados:

/cache/flux_tarot/
├── 384x640/
│   ├── <hash1>.pt
│   └── ...
├── metadata_shard_0000.json
├── metadata.json
└── _hf_dataset/
    └── images/

2. Lanzar entrenamiento con el YAML existente de FLUX

Usa examples/diffusion/finetune/flux_t2i_flow.yaml directamente. El YAML ya selecciona FLUX.1-dev, ajuste fino de transformador completo, el adaptador de flow matching de FLUX, un tamaño de lote efectivo de 32 y FSDP2 con 8 vías.

Proporciona las rutas y configuraciones específicas del tarot como anulaciones de línea de comandos:

uv run --locked --no-default-groups --extra diffusion \
  torchrun --nproc-per-node=8 \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/flux_t2i_flow.yaml \
  --model.transformer_engine_fp8 false \
  --data.dataloader.cache_dir /cache/flux_tarot \
  --data.dataloader.base_resolution '[384,640]' \
  --lr_scheduler.lr_decay_style constant \
  --lr_scheduler.lr_warmup_steps 20 \
  --step_scheduler.max_steps 200 \
  --step_scheduler.ckpt_every_steps 50 \
  --checkpoint.checkpoint_dir /tmp/flux_tarot/checkpoints/full \
  --checkpoint.save_consolidated true \
  --seed 2026

La ejecución produce checkpoints en los pasos 50, 100, 150 y 200. El checkpoint final se etiqueta epoch_66_step_199; la etiqueta es basada en cero aunque representa el paso 200 del optimizador completado.

3. Generar desde el checkpoint ajustado fino

Usa el YAML de generación de FLUX existente y apunta model.checkpoint al checkpoint de entrenamiento completo:

uv run --locked --no-default-groups --extra diffusion \
  python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint /tmp/flux_tarot/checkpoints/full/epoch_66_step_199 \
  --inference.height 640 \
  --inference.width 384 \
  --inference.prompts '["a trtcrd of an astronaut tending a rose garden on Mars, \"the gardener\""]' \
  --output.output_dir /tmp/flux_tarot/generations/full/step_200 \
  --seed 2026

Incluye trtcrd para invocar el estilo de tarot aprendido. Para una comparación de control, mantén la semilla y la escena fijas pero omite el disparador:

uv run --locked --no-default-groups --extra diffusion \
  python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint /tmp/flux_tarot/checkpoints/full/epoch_66_step_199 \
  --inference.height 640 \
  --inference.width 384 \
  --inference.prompts '["an astronaut tending a rose garden on Mars, \"the gardener\""]' \
  --output.output_dir /tmp/flux_tarot/generations/control \
  --seed 2026

Resultados

En el paso 200, los prompts de astronauta con disparador conservan su contenido solicitado mientras adquieren una paleta vintage crema, roja y negra, contornos de tinta gruesos, campos de color planos, tonos de papel envejecido y composición alegórica de cartas. El astronauta sin disparador permanece fotográfico, demostrando que el efecto aprendido está sustancialmente asociado con trtcrd en lugar de reemplazar el modelo base globalmente.

Prompt (semilla)Línea baseAjustado fino (paso 200)Astronauta cuidando un jardín de rosas en Marte (semilla 2026)

4. Rendimiento

Todas las mediciones se recogieron en un nodo con 8 GPUs NVIDIA H100 80GB. Los resultados son medias ± desviación estándar de la muestra sobre tres ventanas de 10 pasos en estado estacionario.

Texto a imagen — 512×512

ModeloEntrenamientoParalelismoGBS / LBSTiempo por pasoImágenes/sImágenes/s/GPUPico asignado/GPUFLUX.1-devCompletoFSDP232 / 40.902 ± 0.039 s35.51 ± 1.554.44 ± 0.1963.88 GiBFLUX.1-devLoRA r64DDP48 / 60.894 ± 0.008 s53.73 ± 0.486.72 ± 0.0667.43 GiBQwen-ImageCompletoFSDP240 / 50.974 ± 0.075 s41.21 ± 3.065.15 ± 0.3853.55 GiBQwen-ImageLoRA r64DDP24 / 30.515 ± 0.006 s46.63 ± 0.545.83 ± 0.0766.33 GiB

Texto a video — 512×512×49 fotogramas

Cada muestra es un clip de video de 49 fotogramas.

ModeloEntrenamientoGBS / LBSCheckpointing de activaciónTiempo por pasoClips/sClips/s/GPUPico asignado/GPUWan 2.1 1.3BCompleto8 / 1Apagado0.942 ± 0.038 s8.50 ± 0.351.06 ± 0.046.09 GiBWan 2.1 14BCompleto8 / 1Encendido3.798 ± 0.017 s2.107 ± 0.0060.263 ± 0.00633.35 GiBWan 2.1 14BLoRA r6416 / 2Encendido7.585 ± 0.014 s2.110 ± 0.0000.263 ± 0.00024.07 GiBWan 2.2 A14B, alto ruidoCompleto8 / 1Encendido4.628 ± 0.031 s1.730 ± 0.0100.217 ± 0.00623.57 GiBHunyuanVideo 1.5Completo8 / 1Encendido5.926 ± 0.046 s1.350 ± 0.0100.170 ± 0.00015.90 GiBHunyuanVideo 1.5LoRA r648 / 1Encendido5.575 ± 0.006 s1.433 ± 0.0060.180 ± 0.00010.58 GiB

Detalles de medición

• Hardware: 8× H100 80GB HBM3, completamente conectados por NVLink.

• Conjunto de datos de imagen: lambda/naruto-blip-captions, 256 muestras en caché.

• Conjunto de datos de video: svjack/Lelouch_Vi_Britannia_FramePack_First_Last_Frame_Video_Captioned, 112 muestras en caché.

• Se aplicaron lotes completos con drop_last=true.

• Las escrituras de checkpoint se desactivaron.

• El tiempo por paso incluye carga de datos, forward, backward, recorte de gradientes, optimizador y trabajo del scheduler.

• La memoria es el uso máximo del asignador CUDA de PyTorch, no la memoria total del dispositivo NVML.

Otros ejemplos de ajuste fino/LoRA

Los resultados del ajuste fino y LoRA muestran el poder de NeMo Automodel para la especialización en dominios. Por ejemplo, ajustar fino el modelo Wan 2.1 en un conjunto de datos de video Ghibli adaptó con éxito el estilo de salida, demostrado por un cambio notable en la apariencia de una flor en comparación con la línea base.

Línea base:

Ajustado fino en videos de Ghibli:

También observamos el impacto distintivo de usar LoRA, donde aplicar el adaptador a Wan 2.1 hizo que el video adoptara un estilo Ghibli característico, particularmente visible en el resaltado de los ojos de los personajes.

Sin LoRA:

Con LoRA:

Estos ejemplos, incluyendo los de FLUX.2, confirman que los usuarios pueden lograr tanto la máxima calidad mediante ajuste fino completo como la máxima eficiencia mediante PEFT estilo LoRA, adaptando la salida a dominios estilísticos específicos.

Pruébalo hoy

Aprende más sobre la integración y encuentra más ejemplos de ajuste fino en la documentación de NeMo Automodel

Próximamente: APIs de recetas en Python

YAML es un buen ajuste para la configuración reproducible, especialmente para equipos que quieren archivos que puedan verificar, revisar y reutilizar, pero muchos equipos también necesitan una interfaz programática.

En una próxima versión de NeMo Automodel, planeamos exponer las recetas de difusión también a través de una API Pythonica completamente tipada. Los usuarios podrán componer el mismo modelo, datos, optimizador, PEFT/LoRA, paralelismo, checkpointing y piezas de generación directamente desde Python.

El camino Pythonico está destinado a hacer que las recetas sean más fáciles de usar desde el código de entrenamiento existente, notebooks y flujos de trabajo de experimentos, y ofrecer una interfaz Pythonica de primera clase junto con la ruta rápida de YAML.