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

使用 NVIDIA NeMo Automodel 和 🤗 Diffusers 扩展视频和图像模型的微调能力

返回文章列表

[0

[-1

企业级

+

]

文章

发布于 2026 年 7 月 17 日

点赞

41

[

• +35

Pranav Prashant Thombre

pthombre

关注

nvidia

linnan wang

linnanwang

Alexandros Koumparoulis

akoumpa

Wenwen Gao

wgao2021

Sylendran Arunagiri

Sylendran95

Bernard Nguyen

bernardwin

NVIDIA 与 Hugging Face 联合发布。特别感谢 Hugging Face 的 Sayak Paul 在集成工作和本文合著方面的贡献。

扩散模型推动了过去两年一些最令人兴奋的开源发布——例如用于文生图的 FLUX.1-dev，以及用于文生视频的 Wan 2.1 和 HunyuanVideo。🤗 Diffusers 库已成为这些模型的默认家园，为研究人员和开发者提供了一个统一且一致的接口，用于推理、适配和流水线组合。

此外，扩散模型的训练和微调需求也在上升，这需要提供内存高效的分片、潜在缓存、多分辨率分桶以及从单个 GPU 到数百个 GPU 的可扩展配置的工具。

为满足这些技术需求，我们推出了开源的 NVIDIA NeMo Automodel 库。今天，我们重点介绍 NVIDIA 与 Hugging Face 的合作，该合作将生产级分布式扩散训练能力扩展到 Hugging Face Hub 上的任何 Diffusers 格式模型——无需检查点转换，也无需对任何新模型进行重写。该集成已在 Diffusers 训练指南中进行了文档说明，并完全开源，遵循 Apache 2.0 协议。

目录

• 什么是 NeMo Automodel？
• 支持的扩散模型
• 本次合作带来的突破
• 微调工作流程概览 1. 预编码数据集 2. 使用现有 FLUX YAML 启动训练 3. 从微调检查点生成 4. 性能表现
• 其他微调/LoRA 示例
• 立即体验
• 下一步：Pythonic 食谱 API
• 资源

什么是 NeMo Automodel？

NeMo Automodel 是一个开源的 PyTorch DTensor 原生训练库，属于 NVIDIA NeMo 框架的一部分，围绕 Diffusers 生态系统中两个关键设计原则构建：

• Hugging Face 原生支持。将 pretrained_model_name_or_path 指向 Hugging Face Hub 上的任何 Diffusers 模型 ID 即可开始训练。NeMo Automodel 使用 Diffusers 模型类（如 WanTransformer3DModel）进行加载，并使用 Diffusers 流水线（如 WanPipeline）进行生成。检查点可以无缝地重新导入 Diffusers 生态系统。

• 一个程序，任意规模。配方和训练脚本可以轻松修改以适应任何规模的训练。并行性是一个配置选择，而非代码重写——通过声明配置即可在 FSDP2、张量并行、专家并行、上下文并行和流水线并行之间切换，无需重写模型。

目前 AutoModel 仅支持流匹配模型。其内部使用流匹配作为训练目标，通过预编码 VAE 输出的潜在空间训练和多分辨率分桶数据加载来加速吞吐量。

支持的扩散模型

NeMo Automodel 集成包含以下开源扩散模型的即用型微调配方。列表反映了 examples/diffusion/finetune 中当前的配方。

模型

Hugging Face ID

任务

参数

LoRA 配方

Wan 2.1 T2V 1.3B / 14B

Wan-AI/Wan2.1-T2V-1.3B-Diffusers

Wan-AI/Wan2.1-T2V-14B-Diffusers

/

文本到视频

1.3B（可适配单块40GB A100）/ 14B

是

Wan 2.2 T2V A14B

Wan-AI/Wan2.2-T2V-A14B-Diffusers

27B总参数量（MoE），每步激活14B

否

FLUX.1-dev

black-forest-labs/FLUX.1-dev

文本到图像

12B

FLUX.2-dev

black-forest-labs/FLUX.2-dev

32B

HunyuanVideo 1.5

hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v

13B

Qwen-Image

Qwen/Qwen-Image

20B（MMDiT）

此次合作带来的突破

对于Diffusers用户而言，实际收益体现在几个具体能力上。

无需检查点转换。Hub中的预训练权重可直接使用，无需额外转换为"训练格式"再转换回来。您微调后的检查点可直接加载到DiffusionPipeline进行推理，或重新上传至Hub进行共享。下游工具——量化、编译、LoRA适配器、自定义采样器——均可继续使用。

快速支持新模型。当新扩散模型在Diffusers中上线时，通过NeMo Automodel启用该模型只需添加少量封装代码——数据预处理处理器和模型适配器，而非完整的自定义训练脚本。其余技术栈（FSDP2、分块数据加载、检查点、生成）保持不变，相同的YAML驱动工作流同样适用。

完整且参数高效的微调。支持完整微调和LoRA风格PEFT两种方式，您可在最大质量（大型集群完整微调）和最大效率（单节点LoRA）之间自由选择。相同的配方结构可兼容两种方式。

超越内置脚本的可扩展训练。NeMo Automodel新增了FSDP2、张量、上下文和流水线并行等分片方案，支持多节点编排（当前SLURM，即将支持Kubernetes）和多分辨率分块。这些能力使得训练FLUX.1-dev（12B）和HunyuanVideo（13B）等大型模型成为可能。

微调工作流概览

在本节中，我们将逐步演示如何对支持的任何模型进行微调。推荐的安装方式是使用NeMo Automodel Docker容器（nvcr.io/nvidia/nemo-automodel:26.06），该容器预装了PyTorch、TransformerEngine和其他CUDA编译依赖项。或者通过pip3 install nemo-automodel安装，或从源码安装（pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git）；详见安装指南获取所有选项。

本指南将演示如何在78张显卡的Rider–Waite塔罗数据集上对FLUX.1-dev进行完整Transformer微调，然后从生成的检查点进行生成。它复用已提交的YAML配置，并通过命令行覆盖应用特定运行设置，因此无需创建新配置文件。

1. 预编码数据集

扩散配方使用缓存的VAE潜在表示和文本嵌入，而非在每次训练步骤中对源图像进行编码。直接从Hugging Face流式传输78张Rider–Waite图像，并在所有可见GPU上分发预处理：

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 captions already contain the trtcrd trigger token. With this pixel budget and the dataset's portrait aspect ratio, preprocessing assigns the samples to the 384×640 bucket used by the showcase run.

For image training, preprocessing produces .pt cache files and sharded metadata:

/cache/flux_tarot/
├── 384x640/
│   ├── <hash1>.pt
│   └── ...
├── metadata_shard_0000.json
├── metadata.json
└── _hf_dataset/
    └── images/

2. 使用现有 FLUX YAML 启动训练

直接使用 examples/diffusion/finetune/flux_t2i_flow.yaml 文件。该 YAML 文件已配置 FLUX.1-dev、完整的变压器微调、FLUX flow-matching 适配器、有效的批量大小为 32，以及八路 FSDP2。

通过命令行覆盖提供塔罗牌特定的路径和设置：

uv run --locked --no-default-groups --extra diffusion \
  torchrun --nproc-per-node=8 \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/flux_t2i_flow.yaml \
  --model.transformer_engine_fp8
false
\
  --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
true
\
  --seed 2026

该运行会在第 50、100、150 和 200 步生成检查点。最终检查点标记为 epoch_66_step_199；尽管它表示已完成第 200 次优化步骤，但标签仍从零开始计数。

3. 从微调检查点生成图像

使用现有的 FLUX 生成 YAML，并将 model.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
'["a trtcrd of an astronaut tending a rose garden on Mars, \"the gardener\""]'
\
  --output.output_dir /tmp/flux_tarot/generations/full/step_200 \
  --seed 2026

包含 trtcrd 可以调用学习到的塔罗牌风格。为了进行对照比较，保持种子和场景不变但省略触发词：

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
'["an astronaut tending a rose garden on Mars, \"the gardener\""]'
\
  --output.output_dir /tmp/flux_tarot/generations/control \
  --seed 2026

#### 结果

在第 200 步时，触发的宇航员提示保留了请求的内容，同时获得了奶油色、红色和黑色的复古配色、厚重的墨水轮廓、平铺色块、泛黄纸张色调以及寓言卡的构图。未触发的宇航员保持摄影风格，证明学习到的效果与 trtcrd 显著相关，而非全局替换基础模型。

提示（种子）

基线

微调（第 200 步）

火星上的宇航员照料玫瑰花园（种子 2026）

火星月光下的宇航员花园（种子 2028）

火星下双月的宇航员种植幼苗（种子 2029）

4. 性能

所有测量数据均在一个配备 8× NVIDIA H100 80GB GPU 的节点上采集。结果为三个稳态 10 步窗口的平均值 ± 样本标准差。

#### 文本到图像 — 512×512（查看秒数是否对齐）

训练

并行性

GBS / LBS

步骤时间

图像/秒

图像/秒/GPU

每个 GPU 峰值分配

完整

FSDP2

32 / 4

0.902 ± 0.039 秒

35.51 ± 1.55

4.44 ± 0.19

63.88 GiB

LoRA r64

DDP

48 / 6

0.894 ± 0.008 秒

53.73 ± 0.48

6.72 ± 0.06

67.43 GiB

40 / 5

0.974 ± 0.075 秒

41.21 ± 3.06

5.15 ± 0.38

53.55 GiB

24 / 3

0.515 ± 0.006 秒

46.63 ± 0.54

5.83 ± 0.07

66.33 GiB

#### 文本到视频 — 512×512×49 帧

每个样本为一个 49 帧的视频片段。

激活检查点

片段/秒

片段/秒/GPU

Wan 2.1 1.3B

8 / 1

关闭

0.942 ± 0.038 秒

8.50 ± 0.35

1.06 ± 0.04

6.09 GiB

Wan 2.1 14B

开启

3.798 ± 0.017 秒

2.107 ± 0.006

0.263 ± 0.006

33.35 GiB

16 / 2

7.585 ± 0.014 秒

2.110 ± 0.000

0.263 ± 0.000

24.07 GiB

Wan 2.2 A14B，高噪声

4.628 ± 0.031 秒

1.730 ± 0.010

0.217 ± 0.006

23.57 GiB

5.926 ± 0.046 秒

1.350 ± 0.010

0.170 ± 0.000

15.90 GiB

5.575 ± 0.006 秒

1.433 ± 0.006

0.180 ± 0.000

10.58 GiB

测量细节

• 硬件：8× H100 80GB HBM3，完全通过 NVLink 连接。
• 图像数据集：lambda/naruto-blip-captions，256 个缓存样本。
• 视频数据集：svjack/Lelouch_Vi_Britannia_FramePack_First_Last_Frame_Video_Captioned，112 个缓存样本。
• 强制使用完整批次（drop_last=true）。
• 禁用检查点写入。
• 步骤时间包含数据加载、前向传播、反向传播、梯度裁剪、优化器和调度器工作。
• 内存为 PyTorch CUDA 分配器的峰值使用量，而非总 NVML 设备内存。

其他微调/LoRA 示例

微调和 LoRA 的结果展示了 NeMo Automodel 在领域专业化方面的强大能力。例如，在 Ghibli 视频数据集上微调 Wan 2.1 模型，成功调整了输出风格，与基线相比，花朵的外观有明显变化。

基线：

在 Ghibli 视频上微调后：

我们还观察到使用 LoRA 的显著影响，将适配器应用于 Wan 2.1 使视频采用了典型的 Ghibli 风格，尤其体现在对角色眼睛的突出表现上。

无 LoRA：

LoRA：

这些示例（包括 FLUX.2 的示例）表明，用户可以通过完整微调实现最高质量，也可以通过 LoRA 风格的 PEFT 实现最高效率，从而将输出定制到特定的风格领域。

立即体验

了解更多关于集成的信息，并在 NeMo Automodel 文档中查找更多微调示例。

即将推出：Pythonic 食谱 API

YAML 非常适合可重复的配置，特别是对于需要可提交、审查和重用文件的团队，但许多团队也需要编程接口。

在即将推出的 NeMo Automodel 版本中，我们计划通过完全类型化的 Pythonic API 提供扩散食谱。用户将能够直接从 Python 组合相同的模型、数据、优化器、PEFT/LoRA、并行性、检查点和生成组件。

Pythonic 路径旨在使食谱更容易与现有训练代码、笔记本和实验工作流程集成，并与 YAML 快速启动路径并行提供一流的 Pythonic 接口。

资源

• NeMo Automodel GitHub
• 扩散微调指南
• 扩散数据集准备
• NeMo Automodel for Transformers（LLM/VLM 微调）

本文中提到的模型 6

本文中提到的数据集 3

该作者的更多文章

NVIDIA Nemotron 3 Embed 在 RTEB 排名第一，推动智能代理检索技术发展

39

2026年7月16日

智能代理的数据

24

2026年7月8日

社区

编辑

预览

通过在文本输入框中拖放、粘贴或点击此处上传图片、音频和视频

轻点或粘贴此处上传图片

评论

· 注册或登录以发表评论

• +29