NVIDIA NeMo Automodel Revolutionizes Diffusers Model Fine-Tuning

Introduction to NeMo Automodel

NVIDIA NeMo Automodel is presented as an innovative open-source library designed to optimize the training of image and video diffusion models. Integrated with PyTorch DTensor, this solution is an integral part of NVIDIA's NeMo framework. It relies on two fundamental principles to seamlessly integrate into the Diffusers ecosystem.

First, NeMo Automodel is designed to be fully compatible with Hugging Face. This means that users only need to specify a pre-trained Diffusers model to start training. Diffusers model classes, such as WanTransformer3DModel, are used for loading, while Diffusers pipelines, like WanPipeline, facilitate generation. This integration allows for a smooth transition of checkpoints within the Diffusers ecosystem.

Next, the scale of training is highly adaptable. Users can adjust recipes and training scripts to meet different scales without the need for code rewriting. Parallelism thus becomes a simple configuration, allowing for easy switching between FSDP2, tensor, contextual, and pipeline parallelism.

Supported Diffusion Models

NeMo Automodel offers ready-to-use recipes for fine-tuning various diffusion models. Among these are models such as Wan 2.1 T2V with 1.3B and 14B versions, as well as Wan 2.2 T2V A14B. Other models include development versions of FLUX by black-forest-labs, as well as HunyuanVideo 1.5, optimized for a resolution of 720p.

These models are accessible via specific identifiers on the Hub, facilitating their use for fine-tuning tasks.

Benefits of Collaboration

The integration of NeMo Automodel with Diffusers offers several practical advantages. One of the most notable is the elimination of the need to convert checkpoints. Pre-trained weights can be used directly, without prior conversion. This significantly simplifies the inference and sharing process on the Hub.

Moreover, adding new diffusion models to Diffusers is made easier. Users only need to add a data preprocessing handler and a model adapter, without requiring a complete training script. This approach reduces the time and effort needed to integrate new models.

Fine-tuning is also optimized to be both comprehensive and efficient. Users can choose between full fine-tuning on a large cluster for maximum quality or a PEFT style like LoRA on a single node for maximum efficiency.

Fine-Tuning Workflow

The fine-tuning process with NeMo Automodel follows a structured workflow. The recommended installation is done via the NeMo Automodel Docker container, which includes PyTorch and TransformerEngine optimized for CUDA. Alternatively, installation via pip or from source is possible.

The typical workflow begins with pre-encoding the dataset. Cached VAE latents and text embeddings are used, thus avoiding the need to encode source images at each training step.

Next, training is initiated using an existing YAML file, such as flux_t2i_flow.yaml. This file automatically configures the FLUX.1-dev model for complete fine-tuning, with parameters optimized for efficient batch size and FSDP2 parallelism.

Finally, generation from the fine-tuned checkpoint is carried out using an existing generation YAML, pointing to the complete training checkpoint.

The results obtained demonstrate that triggered astronaut prompts retain their content while adopting a unique visual style, whereas untriggered astronauts remain true to their photographic appearance. This shows that the learned effect is primarily related to the trigger used.