NVIDIA联合Hugging Face推出大规模视频图像模型微调工具

背景

过去两年中，扩散模型（Diffusion Models）驱动了多项最令人兴奋的开源发布，例如用于文生图的 FLUX.1-dev，以及用于文生视频的 Wan 2.1 和 HunyuanVideo。🤗 Diffusers 库已成为这些模型的事实标准，为研究者和开发者提供了统一的推理、适配和流水线组合接口。与此同时，扩散模型的训练和微调需求也在快速增长，这需要具备内存高效分片、潜空间缓存、多分辨率分桶（multiresolution bucketing）以及能从单 GPU 优雅扩展到数百 GPU 的配置工具。

为了满足这些技术需求，NVIDIA 与 Hugging Face 合作推出了 NVIDIA NeMo Automodel 开源库，将生产级、分布式的扩散训练能力带到 Hugging Face Hub 上的任何 Diffusers 格式模型，且无需任何检查点转换或模型重写。该集成已在 Diffusers 训练指南中记录，并完全以 Apache 2.0 许可证开源。

核心内容

什么是 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、张量并行、专家并行、上下文并行和流水线并行之间切换。

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

支持的扩散模型

NeMo Automodel 集成为以下开源扩散模型提供了开箱即用的微调配方（位于 examples/diffusion/finetune 目录下）：

• FLUX.1-dev（文生图）
• Wan 2.1（文生视频）
• HunyuanVideo（文生视频）
• 以及更多将在未来添加的模型。

此次协作解锁的能力

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

• 无需检查点转换：来自 Hub 的预训练权重可直接使用。无需转换到专门的“训练格式”再转换回来。微调后的检查点可直接加载到 DiffusionPipeline 中进行推理，或上传回 Hub 分享。下游工具（量化、编译、LoRA 适配器、自定义采样器）全部保持正常工作。

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

• 完整微调与参数高效微调：既支持完整微调（full fine-tuning），也支持 LoRA 风格的 PEFT。用户可以在最大质量（在大集群上做完整微调）和最大效率（在单节点上做 LoRA）之间做出选择。同一配方结构可处理两者。

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

微调工作流程示例

以下以 FLUX.1-dev 在 78 张 Rider–Waite 塔罗牌数据集上的完整 Transformer 微调为例，展示典型工作流程。

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/
│   ├── <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、完整 Transformer 微调、FLUX 流匹配适配器、有效批次大小 32，以及 8 路 FSDP2。

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

uv run --locked --no-default-groups --extra diffusion \
  torchrun --nproc-per-node=8 \
  examples/diffusion/finetune/finetune.py \
  -