Fine-tune a escala modelos Diffusers con NeMo Automodel

Hace poco NVIDIA y Hugging Face anunciaron una integración que te permite afinar modelos de difusión a escala sin conversiones de checkpoints ni reescrituras de modelo. Suena técnico, pero piensa en esto como una ruta directa: tomas un modelo en el Hub de Diffusers, lo apuntas y empiezas a entrenar, desde una sola GPU hasta clústeres con cientos de GPUs.

Qué es NeMo Automodel

NeMo Automodel es una librería open source de NVIDIA, construida sobre PyTorch y DTensor, diseñada para entrenamiento distribuido de modelos de difusión. Dos principios guían su diseño:

• Nativo Hugging Face. Apunta pretrained_model_name_or_path a cualquier ID de Diffusers en el Hub y comienza a entrenar. NeMo usa las clases y pipelines de Diffusers para cargar y generar, así que los checkpoints son compatibles de ida y vuelta.

• Una sola base, cualquier escala. La paralelización (FSDP2, tensor parallel, pipeline, expert, context, etc.) es una opción de configuración, no una reescritura del código. Eso facilita probar la misma receta en entornos de 1 GPU hasta clusters multi-nodo.

Bajo el capó, Automodel usa flow matching como objetivo de entrenamiento y entrenamiento en espacio latente (latents pre-encodificados por VAE) con bucketed dataloading multiresolución para mejorar el throughput.

Modelos soportados (resumen)

Automodel ya trae recetas listas para varios modelos Diffusers populares. Algunos ejemplos:

• Wan 2.1 T2V (1.3B y 14B) - text-to-video
• Wan 2.2 T2V A14B - text-to-video (MoE)
• FLUX.1-dev y FLUX.2-dev - text-to-image (12B y 32B)
• HunyuanVideo 1.5 - text-to-video (13B)
• Qwen-Image - text-to-image (20B)

La lista completa y las recetas están en examples/diffusion/finetune del repositorio.

Qué desbloquea esta colaboración

• No conversiones de checkpoints. Pesos del Hub funcionan out of the box. El checkpoint fine-tuneado vuelve a cargar en un DiffusionPipeline sin dramas.

• Soporte rápido a nuevos modelos. Añadir un nuevo modelo suele ser una adaptación contenida: un handler de preprocesamiento y un adapter de modelo. El resto de la tubería (FSDP2, bucketed dataloading, checkpointing) funciona igual.

• Fine-tuning completo y PEFT (LoRA). Puedes elegir entre calidad máxima con full fine-tune en grandes clusters o eficiencia con LoRA en una sola máquina.

• Paralelismos y escalabilidad avanzados: FSDP2, tensor, pipeline, context, expert parallel, orquestación multi-nodo (SLURM hoy, Kubernetes próximamente) y bucketed dataloading multiresolución.

Resultado práctico: modelizar estilos específicos (tarot, Ghibli, etc.) o adaptar video y imagenes a dominios con menos fricción técnica.

Flujo de trabajo de fine-tuning (ejemplo técnico)

Instalación recomendada: usar el contenedor Docker oficial nvcr.io/nvidia/nemo-automodel:26.06 que incluye PyTorch y dependencias CUDA. Alternativas:

• pip3 install nemo-automodel
• pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git

Ejemplo conciso: fine-tune de FLUX.1-dev sobre un dataset de 78 cartas Rider–Waite (preprocesado en latents).

1. Preprocesamiento distribuido (streaming desde Hugging Face y creación de caches de VAE latents):

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

2. Ejecución del entrenamiento usando la receta YAML incluida:

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 \
  --data.dataloader.cache_dir /cache/flux_tarot \
  --data.dataloader.base_resolution '[384,640]' \
  --step_scheduler.max_steps 200 \
  --checkpoint.checkpoint_dir /tmp/flux_tarot/checkpoints/full \
  --checkpoint.save_consolidated true \
  --seed 2026

3. Generación con el checkpoint final apuntado en la configuración de generación:

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

La receta demuestra además cómo usar tokens trigger (por ejemplo trtcrd) para invocar un estilo aprendido sin cambiar el comportamiento base del modelo.

Rendimiento y métricas clave

Las mediciones reportadas se ejecutaron en una sola máquina con 8× H100 80GB NVLink. Algunos números resumidos:

• FLUX.1-dev full (FSDP2): ~35 imágenes/s en nodo de 8 GPUs, peak por GPU ~64 GiB.
• FLUX.1-dev LoRA r64 (DDP): ~54 imágenes/s, peak por GPU ~67 GiB, mayor images/s/GPU.
• Wan 2.1 1.3B video: ~8.5 clips/s, memoria por GPU ~6 GiB.
• Wan 2.1 14B full: step time mayor y ~33 GiB pico por GPU con activation checkpointing.

Detalles importantes: step time incluye dataloading, forward/backward, clipping, optimizador y scheduler. Memoria reportada corresponde al peak del allocator de PyTorch CUDA.

Ejemplos de dominio y LoRA

• Fine-tuning completo puede transformar la estética global de un modelo para un dominio específico.
• LoRA ofrece una vía eficiente: adaptación significativa con footprint pequeño. Ejemplo: adaptar Wan 2.1 a un estilo tipo Ghibli mostró cambios visibles en color y ojos de personajes manteniendo eficiencia.

Pythonic API y próximos pasos

Si prefieres programar en lugar de YAML, pronto habrá una API Python totalmente tipada que permitirá componer modelos, datos, optimizadores, PEFT, paralelismo, checkpoints y generación desde código. Esto facilita integrar recetas en notebooks, pipelines de experimentos y sistemas de entrenamiento más complejos.

Entrar al ecosistema NeMo Automodel significa menos fricción para pasar de idea a experimento en imagen y video generativo a gran escala. ¿Quieres adaptar un modelo grande a un nicho visual o probar LoRA en una sola máquina? Aquí tienes la infraestructura lista para ello.

Fuente original

https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel