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

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

使用 NVIDIA NeMo Automodel 和 🤗 Diffusers 大规模微调视频与图像模型

A joint post from NVIDIA and Hugging Face. Special thanks to Sayak Paul from Hugging Face for their contributions to the integration work and for co-authoring this blog. 这是 NVIDIA 与 Hugging Face 的联合发布。特别感谢 Hugging Face 的 Sayak Paul 对集成工作所做的贡献,并共同撰写了这篇博客。

Diffusion models power some of the most exciting open-source releases of the last two years — such as FLUX.1-dev for text-to-image and Wan 2.1 and HunyuanVideo for text-to-video. The 🤗 Diffusers library has become the de facto home for these models, giving researchers and builders a single, consistent interface for inference, adaptation, and pipeline composition. 扩散模型驱动了过去两年中最令人兴奋的一些开源发布,例如用于文生图的 FLUX.1-dev,以及用于文生视频的 Wan 2.1 和 HunyuanVideo。🤗 Diffusers 库已成为这些模型的实际大本营,为研究人员和开发者提供了一个统一、一致的接口,用于推理、适配和流水线组合。

In addition, training and fine-tuning diffusion models are also on the rise, requiring utilities that offer memory-efficient sharding, latent caching, multiresolution bucketing, and configurations that scale gracefully from one GPU to hundreds. To cater to these technical demands, we offer the NVIDIA NeMo Automodel open-source library. 此外,扩散模型的训练和微调需求也在不断增长,这需要能够提供内存高效分片、潜在空间缓存、多分辨率分桶以及能够从单 GPU 平滑扩展到数百 GPU 的配置工具。为了满足这些技术需求,我们推出了 NVIDIA NeMo Automodel 开源库。

Today, we’re highlighting the collaboration between NVIDIA and Hugging Face that brings production-grade, distributed diffusion training to any Diffusers-format model on the Hugging Face Hub — with no checkpoint conversion and no model rewrites for any new model. The integration is documented in the Diffusers training guide and is fully open source under Apache 2.0. 今天,我们重点介绍 NVIDIA 与 Hugging Face 之间的合作,该合作将生产级的分布式扩散模型训练引入了 Hugging Face Hub 上的任何 Diffusers 格式模型——无需进行检查点转换,也无需为任何新模型重写代码。该集成已记录在 Diffusers 训练指南中,并以 Apache 2.0 协议完全开源。


What is NeMo Automodel?

什么是 NeMo Automodel?

NeMo Automodel is an open-source PyTorch DTensor-native training library, part of the NVIDIA NeMo framework, built around two design principles that matter for the Diffusers ecosystem: NeMo Automodel 是一个开源的、原生支持 PyTorch DTensor 的训练库,属于 NVIDIA NeMo 框架的一部分。它围绕着对 Diffusers 生态系统至关重要的两个设计原则构建:

  • Hugging Face native. Point pretrained_model_name_or_path at any Diffusers model ID on the Hub and start training. NeMo Automodel uses Diffusers model classes (e.g. WanTransformer3DModel) for loading and Diffusers pipelines (WanPipeline) for generation. Checkpoints round-trip cleanly back into the Diffusers ecosystem. Hugging Face 原生支持。 只需将 pretrained_model_name_or_path 指向 Hub 上的任何 Diffusers 模型 ID 即可开始训练。NeMo Automodel 使用 Diffusers 模型类(例如 WanTransformer3DModel)进行加载,并使用 Diffusers 流水线(WanPipeline)进行生成。检查点可以完美地回传到 Diffusers 生态系统中。

  • One program, any scale. The recipes and training scripts can be easily modified to suit training at any scale. Parallelism is a configuration choice, not a code rewrite — switch between FSDP2, tensor parallel, expert parallel, context parallel, and pipeline parallel by declaring configurations, not rewriting models. 一套程序,任意规模。 配方和训练脚本可以轻松修改以适应任何规模的训练。并行化是一种配置选择,而非代码重写——通过声明配置而非重写模型,即可在 FSDP2、张量并行、专家并行、上下文并行和流水线并行之间进行切换。

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. AutoModel 目前仅支持流匹配(flow-matching)模型。在底层,它使用流匹配作为训练目标,并通过潜在空间训练(通过预编码的 VAE 输出)和多分辨率分桶数据加载来加速吞吐量。


Supported diffusion models

支持的扩散模型

NeMo Automodel integration ships with ready-to-use fine-tuning recipes for the open diffusion models below. The list reflects the recipes currently in examples/diffusion/finetune. NeMo Automodel 集成附带了针对以下开源扩散模型的即用型微调配方。该列表反映了当前 examples/diffusion/finetune 中的配方。

ModelHugging Face IDTaskParametersLoRA recipe
Wan 2.1Wan-AI/Wan2.1-T2V-1.3B-Diffusers / 14BText-to-Video1.3B / 14BYes
Wan 2.2Wan-AI/Wan2.2-T2V-A14B-DiffusersText-to-Video27B total (MoE)No
FLUX.1-devblack-forest-labs/FLUX.1-devText-to-Image12BYes
FLUX.2-devblack-forest-labs/FLUX.2-devText-to-Image32BYes
HunyuanVideohunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2vText-to-Video13BYes
Qwen-ImageQwen/Qwen-ImageText-to-Image20B (MMDiT)Yes

What this collaboration unlocks

此次合作带来的价值

For Diffusers users, the practical gains break down into a few concrete capabilities: 对于 Diffusers 用户而言,实际收益可归纳为以下几项具体能力:

  • No checkpoint conversion. Pretrained weights from the Hub work out of the box. There’s no separate “training format” to convert to, then convert back. Your fine-tuned checkpoint loads directly into a DiffusionPipeline for inference, or back to the Hub for sharing. Downstream tools — quantization, compilation, LoRA adapters, custom samplers — all keep working. 无需检查点转换。 来自 Hub 的预训练权重可直接使用。无需转换成单独的“训练格式”后再转回。你微调后的检查点可以直接加载到 DiffusionPipeline 中进行推理,或上传回 Hub 进行分享。下游工具(如量化、编译、LoRA 适配器、自定义采样器)均可正常工作。

  • 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 — a data preprocessing handler and a model adapter — rather than a full custom training script. 支持新模型的快速通道。 当新的扩散模型进入 Diffusers 时,在 NeMo Automodel 中启用它只需进行少量的代码添加(一个数据预处理处理器和一个模型适配器),而无需编写完整的自定义训练脚本。

  • Full and parameter-efficient fine-tuning. Both full fine-tuning and LoRA-style PEFT are supported. 全量微调与参数高效微调。 同时支持全量微调和 LoRA 风格的 PEFT。

  • Scalable training. NeMo Automodel adds sharding schemes such as FSDP2, tensor, context, and pipeline parallelisms, multi-node orchestration, and multiresolution bucketing. These capabilities make training larger models like FLUX.1-dev (12B) and HunyuanVideo (13B) possible. 可扩展的训练。 NeMo Automodel 增加了 FSDP2、张量并行、上下文并行、流水线并行、多节点编排以及多分辨率分桶等分片方案。这些能力使得训练像 FLUX.1-dev (12B) 和 HunyuanVideo (13B) 这样的大型模型成为可能。


A look at the fine-tuning workflow

微调工作流程概览

In this section, we walk through the typical workflow for fine-tuning any of the supported models. The recommended way to install Automodel is the NeMo Automodel Docker container (nvcr.io/nvidia/nemo-automodel:26.06), which ships with PyTorch, TransformerEngine, and other CUDA-compiled dependencies pre-built. 在本节中,我们将介绍微调任何受支持模型的典型工作流程。安装 Automodel 的推荐方式是使用 NeMo Automodel Docker 容器 (nvcr.io/nvidia/nemo-automodel:26.06),其中预装了 PyTorch、TransformerEngine 和其他 CUDA 编译的依赖项。