Recently NVIDIA and Hugging Face announced an integration that lets you fine-tune diffusion models at scale without checkpoint conversions or model rewrites. Sounds technical, but think of this as a direct route: you pick a model from the Diffusers Hub, point to it, and start training — from a single GPU up to clusters with hundreds of GPUs.
What is NeMo Automodel
NeMo Automodel is an open-source NVIDIA library, built on PyTorch and DTensor, designed for distributed training of diffusion models. Two principles guide its design:
-
Native Hugging Face. Point
pretrained_model_name_or_pathto any Diffusers ID on the Hub and start training. NeMo uses the Diffusers classes and pipelines to load and generate, so checkpoints are compatible both ways. -
One base, any scale. Parallelization (FSDP2, tensor parallel, pipeline, expert, context, etc.) is a configuration option, not a code rewrite. That makes it easy to run the same recipe in environments from 1 GPU to multi-node clusters.
Under the hood, Automodel uses flow matching as the training objective and trains in latent space (latents pre-encoded by a VAE) with bucketed multi-resolution dataloading to improve throughput.
Supported models (summary)
Automodel already ships recipes ready for several popular Diffusers models. Some examples:
- Wan 2.1 T2V (1.3B and 14B) - text-to-video
- Wan 2.2 T2V A14B - text-to-video (MoE)
- FLUX.1-dev and FLUX.2-dev - text-to-image (12B and 32B)
- HunyuanVideo 1.5 - text-to-video (13B)
- Qwen-Image - text-to-image (20B)
The full list and the recipes live in examples/diffusion/finetune of the repository.
What this collaboration unlocks
-
No checkpoint conversions. Weights from the Hub work out of the box. The fine-tuned checkpoint loads back into a
DiffusionPipelinewithout drama. -
Fast support for new models. Adding a new model is often a contained adaptation: a preprocessing handler and a model adapter. The rest of the pipeline (FSDP2, bucketed dataloading, checkpointing) works the same.
-
Full fine-tuning and PEFT (LoRA). You can choose maximum quality with full fine-tune on large clusters or efficiency with LoRA on a single machine.
-
Advanced parallelisms and scalability: FSDP2, tensor, pipeline, context, expert parallel, multi-node orchestration (SLURM today, Kubernetes soon), and bucketed multi-resolution dataloading.
Practical result: model specific styles (tarot, Ghibli, etc.) or adapt video and images to domains with less technical friction.
Fine-tuning workflow (technical example)
Recommended installation: use the official Docker container nvcr.io/nvidia/nemo-automodel:26.06 which includes PyTorch and CUDA dependencies. Alternatives:
pip3 install nemo-automodelpip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git
Concise example: fine-tune FLUX.1-dev on a dataset of 78 Rider–Waite tarot cards (preprocessed into latents).
- Distributed preprocessing (streaming from Hugging Face and creating caches of 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
- Run training using the included YAML recipe:
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
- Generation with the final checkpoint pointed in the generation config:
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
The recipe also shows how to use trigger tokens (for example trtcrd) to invoke a learned style without changing the base behavior of the model.
Performance and key metrics
The reported measurements were run on a single machine with 8× H100 80GB NVLink. Some summarized numbers:
- FLUX.1-dev full (FSDP2): ~35 images/s on an 8-GPU node, peak per GPU ~64 GiB.
- FLUX.1-dev LoRA r64 (DDP): ~54 images/s, peak per GPU ~67 GiB, higher images/s per GPU.
- Wan 2.1 1.3B video: ~8.5 clips/s, memory per GPU ~6 GiB.
- Wan 2.1 14B full: higher step time and ~33 GiB peak per GPU with activation checkpointing.
Important details: step time includes dataloading, forward/backward, clipping, optimizer and scheduler. Reported memory corresponds to the peak of the PyTorch CUDA allocator.
Domain examples and LoRA
- Full fine-tuning can transform the global aesthetic of a model for a specific domain.
- LoRA offers an efficient path: significant adaptation with a small footprint. Example: adapting Wan 2.1 to a Ghibli-like style showed visible changes in color and character eyes while keeping efficiency.
Pythonic API and next steps
If you prefer coding instead of YAML, there will soon be a fully typed Python API that lets you compose models, data, optimizers, PEFT, parallelism, checkpoints and generation from code. That makes it easier to integrate recipes into notebooks, experiment pipelines and more complex training systems.
Joining the NeMo Automodel ecosystem means less friction to go from idea to experiment in large-scale image and video generative work. Want to adapt a large model to a visual niche or try LoRA on a single machine? The infrastructure is ready for that.
Original source
https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel