Co-posted by NVIDIA and Hugging Face. Special thanks to Sayak Paul from Hugging Face for contributing to the integration effort and co-authoring this blog.
The diffusion model powers some of the most exciting open source releases of the past two years, including text-to-image FLUX.1-dev and text-to-video Wan 2.1 and HunyuanVideo. 🤗 The Diffuser library becomes the de facto home for these models, providing researchers and builders with a single, consistent interface for inference, adaptation, and pipeline configuration.
Furthermore, training and fine-tuning of diffusion models is also on the rise, requiring utilities that provide memory-efficient sharding, latent caching, multi-resolution bucketing, and configurations that scale smoothly from one GPU to hundreds of GPUs.
To meet these technical demands, we provide the NVIDIA NeMo Automodel open source library. Today we will focus on the collaboration between NVIDIA and Hugging Face. This provides production-grade, decentralized diffusion training for any diffuser format model on Hugging Face Hub. There is no need to convert checkpoints or rewrite the model for the new model. This integration is documented in the Diffusers training guide and is fully open source in Apache 2.0.
table of contents
What is NeMo Auto Model?
NeMo Automodel is an open source PyTorch DTensor native training library that is part of the NVIDIA NeMo Framework and is built on two design principles that are important to the diffuser ecosystem.
hug face native. Specify pretrained_model_name_or_path to any diffuser model ID on the hub and start training. NeMo Automodel uses a diffuser model class (such as WanTransformer3DModel) for loading and a diffuser pipeline (WanPipeline) for generation. Checkpoints cleanly shuttle back and forth to the diffuser ecosystem. One program for any size. Recipes and training scripts can be easily modified to suit training of any size. Parallelism is a configuration choice, not a code rewrite. To switch between FSDP2, tensor parallelism, expert parallelism, context parallelism, and pipeline parallelism, declare a configuration rather than rewriting the model.
AutoModel currently only supports flow matching models. Under the hood, we use flow matching as the training objective and use latent space training (via pre-encoded VAE output) and multi-resolution bucket dataloading to speed up throughput.
Supported dissemination models
The NeMo Automodel integration ships with ready-to-use fine-tuning recipes for the following open diffusion models: This list reflects the recipes currently located in example/diffusion/finetune.
What you get from this collaboration
For diffuser users, the practical benefits fall into several specific features.
There are no checkpoint conversions. Pre-trained weights from the hub are ready to use. There is no separate “training format” to convert and then convert back. Fine-tuned checkpoints are loaded directly into the DiffusionPipeline for inference or sent back to the hub for sharing. All downstream tools (quantization, compilation, LoRA adapter, custom sampler) will continue to work.
Fast path to support for new models. Once a new diffusion model is added to Diffusers, enabling it in NeMo Automodel requires minor code additions (data preprocessing handlers and model adapters) rather than a full custom training script. The rest of the recipe stack (FSDP2, bucket dataload, checkpointing, generation) carries over unchanged and the same YAML-driven workflow applies.
Complete parameter-efficient fine-tuning. Both full fine-tuning and LoRA-style PEFT are supported, so you can choose between maximum quality (full FT on large clusters) or maximum efficiency (LoRA on a single node). The same recipe structure handles both.
Scalable training beyond what built-in scripts offer. NeMo Automodel adds sharding schemes like FSDP2, tensors, contexts, pipeline parallelism, multi-node orchestration (now SLURM, Kubernetes), and multi-resolution bucketing. These features enable training of large models such as FLUX.1-dev (12B) and HunyuanVideo (13B).
Tweak workflow overview
This section describes a common workflow for fine-tuning supported models. The recommended method to install Automodel is the NeMo Automodel Docker container (nvcr.io/nvidia/nemo-automodel:26.06). It comes with pre-built PyTorch, TransformerEngine, and other CUDA-compiled dependencies. Alternatively, install using pip3 install nemo-automodel or from source (pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git). See the installation guide for all options.
This guide describes the complete transformer fine-tuning of FLUX.1-dev on the 78-card Rider-Waite Tarot dataset and generating the results from checkpoints. Reuses checked-in YAML configurations and applies run-specific settings as command-line overrides, so no new configuration files are required.
1. Pre-encode the dataset
Diffusion recipes consume cached VAE latents and text embeddings instead of encoding the source images at each training step. Stream 78 Rider-Waite images directly from Hugging Face and distribute preprocessing across all visible GPUs.
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
The caption already contains the trtcrd trigger token. Using this pixel budget and the portrait aspect ratio of the dataset, preprocessing assigns the samples to the 384 × 640 buckets used in the showcase run.
For image training, preprocessing produces a .pt cache file and sharded metadata.
/cache/flux_tarot/ ∴── 384×640/ │ §── .pt │ └── … §── metadata_shard_0000.json §── metadata.json └── _hf_dataset/ └── Image/
2. Start training using your existing FLUX YAML
Use example/diffusion/finetune/flux_t2i_flow.yaml directly. YAML has already selected FLUX.1-dev, full transformer fine-tuning, FLUX flow matching adapter, effective batch size 32, and 8-way FSDP2.
Specify Tarot-specific paths and settings as command-line overrides.
uv run –locked –no-default-groups –extra diffusion \ torchrun –nproc-per-node=8 \examples/diffusion/finetune/finetune.py \ -c example/diffusion/finetune/flux_t2i_flow.yaml \ –model.transformer_engine_fp8 error \ –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 truth \ –Seed 2026
This run generates checkpoints at steps 50, 100, 150, and 200. The last checkpoint is labeled epoch_66_step_199. The label starts at 0, even if it represents the 200th completed optimizer step.
3. Generated from fine-tuned checkpoints
Use your existing FLUX-generated YAML and point model.checkpoint at the full training checkpoint.
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 ‘(“Astronaut tending a rose garden on Mars, \”Gardener\””)’ \ –output.output_dir /tmp/flux_tarot/generations/full/step_200 \ –seed 2026
Include trtcrd to invoke the learned tarot style. To compare controls, keep the seed and scene fixed and omit the trigger.
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 ‘(“Astronaut tending a rose garden on Mars, \”Gardener\””)’ \ –output.output_dir /tmp/flux_tarot/generation/control \ –seed 2026
result
At step 200, the triggered astronaut prompt acquires a vintage palette of cream, red, and black, dark ink outlines, flat color fields, old paper tones, and an allegorical card composition while retaining the requested content. The untriggered astronaut remains in the picture, indicating that the learned effects are effectively tied to the trtcrd rather than replacing the base model entirely.
4. Performance
All measurements were collected on a single node with eight NVIDIA H100 80GB GPUs. Results are the mean ± sample standard deviation over three steady-state 10-step windows.
Text to image — 512×512 (see if you can align the seconds)
Model Training Parallelism GBS / LBS Step Time Images/sec Images/sec/GPU Peak Allocation/GPU FLUX.1-dev Full FSDP2 32 / 4 0.902 ± 0.039 s 35.51 ± 1.55 4.44 ± 0.19 63.88 GiB FLUX.1-dev LoRA r64 DDP 48 / 6 0.894 ± 0.008 seconds 53.73 ± 0.48 6.72 ± 0.06 67.43 GiB Qwen Image Full FSDP2 40 / 5 0.974 ± 0.075 seconds 41.21 ± 3.06 5.15 ± 0.38 53.55 GiB Qwen Image LoRA r64 DDP 24 / 3 0.515 ± 0.006 seconds 46.63 ± 0.54 5.83 ± 0.07 66.33 GiB
Text to video — 512x512x49 frames
Each sample is one 49-frame video clip.
Model Training GBS / LBS Activation Checkpoint Step Time Clips/sec Clips/sec/GPU Peak Allocation/GPU Wan 2.1 1.3B Full 8 / 1 Off 0.942 ± 0.038 sec 8.50 ± 0.35 1.06 ± 0.04 6.09 GiB Wan 2.1 14B Full 8 / 1 On 3.798 ± 0.017 seconds 2.107 ± 0.006 0.263 ± 0.006 33.35 GiB Wan 2.1 14B LoRA r64 16 / 2 On 7.585 ± 0.014 seconds 2.110 ± 0.000 0.263 ± 0.000 24.07 GiB Wan 2.2 A14B, high noise Full 8 / 1 on 4.628 ± 0.031 seconds 1.730 ± 0.010 0.217 ± 0.006 23.57 GiB HunyuanVideo 1.5 Full 8 / 1 on 5.926 ± 0.046 seconds 1.350 ± 0.010 0.170 ± 0.000 15.90 GiB HunyuanVideo 1.5 LoRA r64 8 / 1 on 5.575 ± 0.006 seconds 1.433 ± 0.006 0.180 ± 0.000 10.58 GiB
Measurement details
Other Finetuned/LoRA examples
Fine-tuning and LoRA results demonstrate the power of NeMo Automodel for domain specialization. For example, by fine-tuning the Wan 2.1 model on the Ghibli video dataset, we were able to successfully adapt the output style, with noticeable changes in the appearance of flowers compared to the baseline.
Baseline:
Tweak Ghibli videos:
We also observed clear effects when using LoRA. Applying the adapter to Wan 2.1 gave the video a distinctive Ghibli style, especially the highlights in the characters’ eyes.
No LoRA:
LoRA:
These examples, including the FLUX.2 example, confirm that users can achieve both maximum quality with full fine-tuning and maximum efficiency with LoRA-style PEFT that tailors the output to specific style domains.
try it today
To learn more about the integration, check out the NeMo Automodel documentation for detailed tweaking examples.
What’s next: Pythonic recipe API
YAML is great for reproducible configuration, especially for teams who need files that can be checked in, reviewed, and reused, but many teams also want a programmatic interface.
Future NeMo Automodel releases will also expose spreading recipes through a fully typed Pythonic API. Users will be able to create the same model, data, optimizer, PEFT/LoRA, parallelism, checkpointing, and generation parts directly from Python.
The Pythonic path is intended to make it easy to use recipes from existing training code, notebooks, and experimental workflows, and to provide a first-class Pythonic interface along with a YAML quickstart path.
resource
