NVIDIA and Hugging Face have announced a collaboration that brings production-grade distributed diffusion training to any Diffusers-format model on the Hugging Face Hub. The integration uses the NVIDIA NeMo Automodel open-source library and requires no checkpoint conversion or model rewrites for new models. The work is documented in the Diffusers training guide and is fully open source under Apache 2.0.
What is NeMo Automodel?
NeMo Automodel is an open-source PyTorch DTensor-native training library that is part of the NVIDIA NeMo framework. It is built around two design principles that matter for the Diffusers ecosystem.
First, it is Hugging Face native. Users can point the pretrained_model_name_or_path parameter at any Diffusers model ID on the Hub and start training. NeMo Automodel uses Diffusers model classes such as WanTransformer3DModel for loading and Diffusers pipelines like WanPipeline for generation. Checkpoints round-trip cleanly back into the Diffusers ecosystem.
Second, it offers one program that works at any scale. The recipes and training scripts can be easily modified to suit training at any scale. Parallelism is a configuration choice rather than a code rewrite. Users can switch between FSDP2, tensor parallel, expert parallel, context parallel, and pipeline parallel by declaring configurations instead of rewriting models.
AutoModel currently supports flow-matching models only. Under the hood, it uses flow matching as the training objective with latent-space training via pre-encoded VAE outputs and multiresolution bucketed dataloading to accelerate throughput.
Supported diffusion models
The NeMo Automodel integration ships with ready-to-use fine-tuning recipes for several open diffusion models. The list reflects the recipes currently available in examples/diffusion/finetune.
Wan 2.1 T2V 1.3B and 14B models are supported for text-to-video tasks. The 1.3B version fits on a single 40GB A100 GPU. The 14B version is also available. Both have LoRA recipes.
Wan 2.2 T2V A14B is a text-to-video model with 27B total parameters using a mixture of experts architecture, with 14B active per step. It does not have a LoRA recipe.
FLUX.1-dev is a text-to-image model with 12B parameters. It has a LoRA recipe. FLUX.2-dev is a text-to-image model with 32B parameters and also has a LoRA recipe.
HunyuanVideo 1.5 is a text-to-video model with 13B parameters. It has a LoRA recipe.
Qwen-Image is a text-to-image model with 20B parameters using a multimodal diffusion transformer. It has a LoRA recipe.
What this collaboration unlocks
For Diffusers users, the practical gains break down into several concrete capabilities.
There is no checkpoint conversion. Pretrained weights from the Hub work out of the box. There is no separate training format to convert to and then convert back. A fine-tuned checkpoint loads directly into a DiffusionPipeline for inference or back to the Hub for sharing. Downstream tools such as quantization, compilation, LoRA adapters, and custom samplers all keep working.
There is a fast path to new model support. When a new diffusion model lands in Diffusers, enabling it in NeMo Automodel takes a small, contained code addition. This includes a data preprocessing handler and a model adapter rather than a full custom training script. The rest of the recipe stack including FSDP2, bucketed dataloading, checkpointing, and generation carries over unchanged. The same YAML-driven workflow applies.
Both full and parameter-efficient fine-tuning are supported. Users can choose between maximum quality with full fine-tuning on a large cluster or maximum efficiency with LoRA-style PEFT on a single node. The same recipe structure handles both.
Scalable training goes beyond what built-in scripts offer. NeMo Automodel adds sharding schemes such as FSDP2, tensor, context, and pipeline parallelisms. It also includes multi-node orchestration with SLURM today and Kubernetes coming, as well as multiresolution bucketing. These capabilities make training larger models like FLUX.1-dev at 12B parameters and HunyuanVideo at 13B parameters possible.
A look at the fine-tuning workflow
The typical workflow for fine-tuning any of the supported models involves several steps. The recommended way to install Automodel is the NeMo Automodel Docker container at nvcr.io/nvidia/nemo-automodel:26.06, which ships with PyTorch, TransformerEngine, and other CUDA-compiled dependencies pre-built. Alternatively, users can install with pip3 install nemo-automodel or from source with pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git. The installation guide covers all options.
The guide walks through a full-transformer fine-tune of FLUX.1-dev on the 78-card Rider-Waite tarot dataset, then generating from the resulting checkpoint. It reuses the checked-in YAML configs and applies run-specific settings as command-line overrides, so no new config files are required.
1. Pre-encode the dataset
The diffusion recipe consumes cached VAE latents and text embeddings instead of encoding source images during every training step. Users stream the 78 Rider-Waite images directly from Hugging Face and distribute preprocessing across all visible GPUs using a command that specifies the dataset name, media column, caption column, streaming option, maximum images, output directory, processor, model name, and maximum pixels.
The captions already contain the trtcrd trigger token. With this pixel budget and the dataset's portrait aspect ratio, preprocessing assigns the samples to the 384x640 bucket used by the showcase run.
For image training, preprocessing produces .pt cache files and sharded metadata in a directory structure that includes resolution subdirectories, metadata shards, and a dataset folder.
2. Launch training with the existing FLUX YAML
Users can use the examples/diffusion/finetune/flux_t2i_flow.yaml file directly. The YAML already selects FLUX.1-dev, full transformer fine-tuning, the FLUX flow-matching adapter, an effective batch size of 32, and eight-way FSDP2.
Users supply the tarot-specific paths and settings as command-line overrides. The run produces checkpoints at steps 50, 100, 150, and 200. The final checkpoint is labeled epoch_66_step_199, which is zero-based even though it represents the completed 200th optimizer step.
Stay updated
Get the day's AI and automation news in your inbox. No spam, unsubscribe anytime.
3. Generate from the fine-tuned checkpoint
Users can use the existing FLUX generation YAML and point model.checkpoint at the complete training checkpoint. The generation command includes the trtcrd trigger to invoke the learned tarot style. For a control comparison, users can keep the seed and scene fixed but omit the trigger.
At step 200, the triggered astronaut prompts retain their requested content while acquiring a cream, red, and black vintage palette, heavy ink contours, flat color fields, aged-paper tones, and allegorical card composition. The untriggered astronaut remains photographic, demonstrating that the learned effect is substantially associated with trtcrd rather than replacing the base model globally.
4. Performance
All measurements were collected on one node with 8x NVIDIA H100 80GB GPUs. Results are means plus or minus sample standard deviation over three steady-state 10-step windows.
For text-to-image at 512x512 resolution, FLUX.1-dev full training with FSDP2 achieved a step time of 0.902 seconds, 35.51 images per second, 4.44 images per second per GPU, and peak allocated memory of 63.88 GiB per GPU. FLUX.1-dev LoRA with r64 and DDP achieved a step time of 0.894 seconds, 53.73 images per second, 6.72 images per second per GPU, and peak allocated memory of 67.43 GiB per GPU.
Qwen-Image full training with FSDP2 achieved a step time of 0.974 seconds, 41.21 images per second, 5.15 images per second per GPU, and peak allocated memory of 53.55 GiB per GPU. Qwen-Image LoRA with r64 and DDP achieved a step time of 0.515 seconds, 46.63 images per second, 5.83 images per second per GPU, and peak allocated memory of 66.33 GiB per GPU.
For text-to-video at 512x512 resolution with 49 frames, each sample is one 49-frame video clip. Wan 2.1 1.3B full training with a global batch size of 8 and local batch size of 1, with activation checkpointing off, achieved a step time of 0.942 seconds, 8.50 clips per second, 1.06 clips per second per GPU, and peak allocated memory of 6.09 GiB per GPU.
Wan 2.1 14B full training with a global batch size of 8 and local batch size of 1, with activation checkpointing on, achieved a step time of 3.798 seconds, 2.107 clips per second, 0.263 clips per second per GPU, and peak allocated memory of 33.35 GiB per GPU.
Wan 2.1 14B LoRA with r64, a global batch size of 16 and local batch size of 2, with activation checkpointing on, achieved a step time of 7.585 seconds, 2.110 clips per second, 0.263 clips per second per GPU, and peak allocated memory of 24.07 GiB per GPU.
Wan 2.2 A14B high-noise full training with a global batch size of 8 and local batch size of 1, with activation checkpointing on, achieved a step time of 4.628 seconds, 1.730 clips per second, 0.217 clips per second per GPU, and peak allocated memory of 23.57 GiB per GPU.
HunyuanVideo 1.5 full training with a global batch size of 8 and local batch size of 1, with activation checkpointing on, achieved a step time of 5.926 seconds, 1.350 clips per second, 0.170 clips per second per GPU, and peak allocated memory of 15.90 GiB per GPU.
HunyuanVideo 1.5 LoRA with r64, a global batch size of 8 and local batch size of 1, with activation checkpointing on, achieved a step time of 5.575 seconds, 1.433 clips per second, 0.180 clips per second per GPU, and peak allocated memory of 10.58 GiB per GPU.
Measurement details include hardware of 8x H100 80GB HBM3 fully NVLink-connected. The image dataset used was lambda/naruto-blip-captions with 256 cached samples. The video dataset used was svjack/Lelouch_Vi_Britannia_FramePack_First_Last_Frame_Video_Captioned with 112 cached samples. Full batches were enforced with drop_last set to true. Checkpoint writes were disabled. Step time includes dataloading, forward, backward, gradient clipping, optimizer, and scheduler work. Memory is peak PyTorch CUDA allocator usage, not total NVML device memory.
Other Finetuned/LoRA examples
The results from fine-tuning and LoRA showcase the power of NeMo Automodel for domain specialization. Fine-tuning the Wan 2.1 model on a Ghibli video dataset successfully adapted the output style, demonstrated by a noticeable change in a flower's appearance compared to the baseline.
The distinct impact of using LoRA was also observed. Applying the adapter to Wan 2.1 caused the video to adopt a characteristic Ghibli style, particularly visible in the highlighting of characters' eyes.
These examples, including those for FLUX.2, confirm that users can achieve both maximum quality via full fine-tuning and maximum efficiency via LoRA-style PEFT, tailoring the output to specific stylistic domains.
Try it today
Users can learn more about the integration and find more fine-tuning examples in the NeMo Automodel documentation.
Coming next: Pythonic recipe APIs
YAML is a strong fit for reproducible configuration, especially for teams that want files they can check in, review, and reuse. However, many teams also need a programmatic interface.
In an upcoming NeMo Automodel release, the team plans to surface the diffusion recipes through a fully typed Pythonic API as well. Users will be able to compose the same model, data, optimizer, PEFT/LoRA, parallelism, checkpointing, and generation pieces directly from Python.
The Pythonic path is intended to make the recipes easier to use from existing training code, notebooks, and experiment workflows. It will offer a first-class Pythonic interface alongside the YAML quick-start path.
Resources
NeMo Automodel GitHub repository is available for source code. The Diffusion Fine-Tuning Guide provides detailed instructions. The Diffusion Dataset Preparation guide covers dataset handling. NeMo Automodel for Transformers covers LLM and VLM fine-tuning.
Models mentioned in this article include Qwen/Qwen-Image for text-to-image with 20B parameters, Wan-AI/Wan2.1-T2V-1.3B-Diffusers and Wan-AI/Wan2.2-T2V-A14B-Diffusers for text-to-video, black-forest-labs/FLUX.1-dev and black-forest-labs/FLUX.2-dev for text-to-image, and hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v for text-to-video.
Datasets mentioned include lambda/naruto-blip-captions, multimodalart/1920-raider-waite-tarot-public-domain, and svjack/Lelouch_Vi_Britannia_FramePack_First_Last_Frame_Video_Captioned.