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

来源：https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel

返回文章列表 (https://huggingface.co/blog)

• 目录 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#table-of-contents)
• 什么是 NeMo Automodel？ (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#what-is-nemo-automodel)
• 支持的扩散模型 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#supported-diffusion-models)
• 此次合作解锁了什么 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#what-this-collaboration-unlocks)
• 微调工作流程一览 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#a-look-at-the-fine-tuning-workflow)
  • 1. 预编码数据集 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#1-pre-encode-the-dataset)
  • 2. 使用现有 FLUX YAML 启动训练 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#2-launch-training-with-the-existing-flux-yaml)
  • 3. 从微调后的检查点生成 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#3-generate-from-the-fine-tuned-checkpoint)
  • 4. 性能 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#4-performance)
• 其他微调 / LoRA 示例 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#other-finetunedlora-examples)
• 立即尝试 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#try-it-today)
• 即将推出：Python 风格的 Recipe API (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#coming-next-pythonic-recipe-apis)
• 资源 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#resources)

NVIDIA 与 Hugging Face 联合发布。特别感谢 Hugging Face 的 Sayak Paul 对集成工作的贡献以及共同撰写本文。

扩散模型驱动了过去两年最激动人心的开源发布——例如用于文本到图像的 FLUX.1-dev (https://huggingface.co/black-forest-labs/FLUX.1-dev)，以及用于文本到视频的 Wan 2.1 (https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers) 和 HunyuanVideo (https://huggingface.co/hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v)。🤗 Diffusers (https://github.com/huggingface/diffusers) 库已成为这些模型的事实归宿，为研究人员和开发者提供了一致统一的推理、适配和流水线组合接口。此外，扩散模型的训练与微调也日益增多，需要能够提供内存高效分片、潜变量缓存、多分辨率分桶以及从单 GPU 到数百 GPU 无缝扩展的配置等实用工具。为了满足这些技术需求，我们推出了 NVIDIA NeMo Automodel (https://github.com/NVIDIA-NeMo/Automodel) 开源库。今天，我们重点介绍 NVIDIA 与 Hugging Face 的合作，将生产级分布式扩散训练引入 Hugging Face Hub 上的任何 Diffusers 格式模型——无需检查点转换，也无需为任何新模型重写模型。该集成已记录在 Diffusers 训练指南 (https://docs.nvidia.com/nemo/automodel/recipes-e2e-examples/diffusion-fine-tuning) 中，并以 Apache 2.0 协议完全开源。

目录

• 什么是 NeMo Automodel？ (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#what-is-nemo-automodel)
• 支持的扩散模型 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#supported-diffusion-models)
• 此次合作解锁了什么 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#what-this-collaboration-unlocks)
• 微调工作流程一览 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#a-look-at-the-fine-tuning-workflow)
  • 1. 预编码数据集 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#1-pre-encode-the-dataset)
  • 2. 使用现有 FLUX YAML 启动训练 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#2-launch-training-with-the-existing-flux-yaml)
  • 3. 从微调后的检查点生成 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#3-generate-from-the-fine-tuned-checkpoint)
  • 4. 性能 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#4-performance)
• 其他微调 / LoRA 示例 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#other-finetunedlora-examples)
• 立即尝试 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#try-it-today)
• 即将推出：Python 风格的 Recipe API (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#coming-next-pythonic-recipe-apis)
• 资源 (https://huggingface.co/blog/nvidia/scale-diffusers-finetuning-nemo-automodel#resources)

什么是 NeMo Automodel？

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

• 原生的 Hugging Face 体验。 只需将 pretrained_model_name_or_path 指向 Hub 上的任意 Diffusers 模型 ID 即可开始训练。NeMo Automodel 使用 Diffusers 模型类（例如 WanTransformer3DModel）进行加载，并使用 Diffusers 流水线（如 WanPipeline）进行生成。检查点可以干净地往返于 Diffusers 生态。

• 同一程序，任意规模。 配方和训练脚本可以轻松修改以适应任何规模的训练。并行性是一个配置选项，而不是代码重写——通过声明配置来切换 FSDP2、Tensor Parallel、Expert Parallel、Context Parallel 和 Pipeline Parallel，无需重写模型。

AutoModel 目前仅支持 flow-matching 模型。内部使用 flow matching 作为训练目标，采用潜变量空间训练（通过预编码的 VAE 输出）和多分辨率分桶数据加载以加速吞吐。

支持的扩散模型

NeMo Automodel 集成附带了针对以下开放扩散模型的即用微调配方。列表反映了当前 examples/diffusion/finetune 中的配方。

此次合作解锁了什么

对于 Diffusers 用户而言，实际收益体现为几项具体能力。

无需检查点转换。 Hub 上的预训练权重开箱即用。无需先转换为单独的“训练格式”再转换回来。微调后的检查点可直接加载到 DiffusionPipeline 中进行推理，或上传回 Hub 进行共享。下游工具——量化、编译、LoRA 适配器、自定义采样器——均保持正常工作。

快速支持新模型。 当 Diffusers 中出现新的扩散模型时，在 NeMo Automodel 中启用它只需少量、自包含的代码添加——一个数据预处理处理器和一个模型适配器——而无需编写完整的自定义训练脚本。其余的配方栈（FSDP2、分桶数据加载、检查点、生成）保持不变，相同的 YAML 驱动工作流程同样适用。

全量微调与参数高效微调。 同时支持全量微调和 LoRA 风格的 PEFT，因此您可以选择最大质量（在大型集群上进行全量微调）或最高效率（在单个节点上进行 LoRA）。相同的配方结构可以处理两者。

扩展训练，超越内置脚本的能力。 NeMo Automodel 增加了分片方案如 FSDP2、Tensor、Context 以及 Pipeline 并行，多节点编排（当前为 SLURM，即将支持 Kubernetes），以及多分辨率分桶。这些能力使得训练像 FLUX.1-dev（12B）和 HunyuanVideo（13B）这样的大型模型成为可能。

微调工作流程一览

在本节中，我们逐步介绍微调任何支持模型的典型工作流程。安装 Automodel 的推荐方式是使用 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）；参见安装指南 (https://docs.nvidia.com/nemo/automodel/latest/get-started/installation) 获取所有选项。

本指南将逐步展示在 78 张 Rider–Waite tarot 数据集 (https://huggingface.co/datasets/multimodalart/1920-raider-waite-tarot-public-domain) 上对 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

标题已经包含 trtcrd 触发词。在此像素预算和数据集的纵向宽高比下，预处理将样本分配到展示运行所使用的 384×640 分桶。对于图像训练，预处理生成 .pt 缓存文件和分片元数据：

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

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

直接使用 examples/diffusion/finetune/flux_t2i_flow.yaml。YAML 已经选择了 FLUX.1-dev、全量 Transformer 微调、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 相关，而不是全局替换了基础模型。

4. 性能

所有测量均在包含 8 块 NVIDIA H100 80GB GPU 的单节点上完成。结果为三个稳态 10 步窗口的均值 ± 样本标准差。

文本到图像 — 512×512

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

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

测量细节

• 硬件：8× H100 80GB HBM3，全部通过 NVLink 连接。
• 图像数据集：lambda/naruto-blip-captions (https://huggingface.co/datasets/lambda/naruto-blip-captions)，256 个缓存样本。
• 视频数据集：svjack/Lelouch_Vi_Britannia_FramePack_First_Last_Frame_Video_Captioned (https://hugging