Fine-tuning Diffusion Models at Scale with NVIDIA NeMo Automodel and Diffusers

NVIDIA NeMo Automodel is an open-source PyTorch DTensor-native training library, part of the NeMo framework, that brings production-grade distributed diffusion training to any Diffusers-format model on the Hugging Face Hub. The library eliminates checkpoint conversion: pretrained weights from the Hub work directly, and fine-tuned checkpoints load into a standard DiffusionPipeline. NeMo Automodel supports flow-matching models only and uses latent-space training with pre-encoded VAE outputs and multiresolution bucketed dataloading.

Supported models include FLUX.1-dev, Wan 2.1, HunyuanVideo, and Qwen-Image, with ready-to-use YAML recipes. The fine-tuning workflow has three steps: (1) pre-encode the dataset to produce cached VAE latents and text embeddings, (2) launch training using an existing YAML config with command-line overrides, and (3) generate from the checkpoint using the same generation pipeline. Both full fine-tuning and LoRA-style PEFT are supported, and parallelism is a configuration choice (FSDP2, tensor, context, pipeline, expert) without code rewrites.

Performance measurements on a single node of 8× NVIDIA H100 80GB GPUs show, for example, FLUX.1-dev full fine-tuning at 512×512 achieves 35.5 images/s with FSDP2 and a global batch size of 32, peak allocated memory 63.88 GiB per GPU. For video, Wan 2.1 14B full fine-tuning with activation checkpointing runs at 2.1 clips/s (49-frame clips) with 33.35 GiB per GPU. LoRA training for Wan 2.1 14B reduces memory to 24.07 GiB per GPU while maintaining throughput.

The collaboration enables fast support for new models: when a new diffusion model lands in Diffusers, enabling it in NeMo Automodel requires only a data preprocessing handler and a model adapter. All existing recipe infrastructure (FSDP2, bucketed dataloading, checkpointing, generation) carries over without modification. The library is Apache 2.0 licensed, and an upcoming release will add a fully typed Pythonic API alongside the current YAML-driven workflow.

Fine-tune video and image models at scale with NVIDIA NeMo Automodel and 🤗 Diffusers

View Original