NVIDIA
Nvidia和Hugging Face联手,让分布式扩散训练变得轻松简单
Nvidia的NeMo Automodel现已直接集成Hugging Face Diffusers,为FLUX、Wan 2.1和HunyuanVideo等模型提供生产级分布式训练。这个基于Apache 2.0许可的库将并行性作为配置开关处理,并让微调后的检查点可直接加载回推理管线。

扩散模型越来越大,训练它们也变得越来越麻烦。FLUX.1-dev有120亿个参数。Wan 2.1的140亿参数变体需要张量并行才能装进一块H100。HunyuanVideo 1.5有130亿参数。这些模型的推理扩展相对成熟,但训练扩展并非如此,尤其是当你想从单GPU的LoRA实验转移到跨集群的全量微调,而又不想重写训练脚本时。
这正是Nvidia和Hugging Face通过NeMo Automodel与Diffusers库的新集成所瞄准的差距。该集成现已记录在Diffusers训练指南中,并以Apache 2.0许可发布。这里的思路很直接:你指向Hub上任何兼容Diffusers的模型,同一个由YAML驱动的工作流就能处理FSDP2分片、张量并行、上下文并行、潜在缓存和多分辨率分桶。无需转换步骤。无需独立的训练格式。检查点以Diffusers格式输出,可直接加载到扩散管线中进行推理。
对于那些一直在为每个新模型发布拼凑自定义训练脚本的团队来说,这种整合比任何单一性能数字都更重要。
实际变化
实际收益可分为三部分。首先,无需检查点转换。微调后的NeMo检查点就是一个Diffusers检查点,相同的文件夹结构,相同的模型类。下游工具(量化、编译、自定义采样器、LoRA合并)无需适配步骤就能继续工作。其次,快速支持新模型。当一个新的扩散架构出现在Diffusers中时,在NeMo Automodel中启用它只需要一个数据预处理处理器和一个模型适配器,而不是一个完整的自定义训练脚本。其余的配方栈(并行、分桶、检查点、生成)保持不变。

第三,全量微调和LoRA风格的PEFT都由相同的配方结构处理。YAML配置选择运行哪一个;代码路径是共享的。这意味着从业者可以先在单节点上使用LoRA验证数据质量,然后通过更改几个配置字段和一个并行性声明,而非替换训练脚本,扩展到跨八块(或八十块)GPU的全量微调。
基准测试层面
在8块H100 80GB上发布的基准测试给出了该库的性能水平。FLUX.1-dev在512x512分辨率下的全量微调以每秒35.5张图像的速度运行,全局批次大小为32,峰值GPU分配为63.9 GiB。使用秩为64且DDP的LoRA,速度提升至每秒53.7张图像,消耗67.4 GiB。对于文本到视频,Wan 2.1的14B模型对49帧512x512片段进行全量微调,开启激活检查点后,大约每秒2.1个片段,峰值每GPU 33.4 GiB。1.3B变体则轻松适配,峰值6.1 GiB,无需激活检查点。
这些数字是禁用检查点写入并强制执行完整批次后的稳态平均值。实际运行中定期检查点会稍慢一些,但单GPU LoRA与分布式全量微调之间的差距足够小,大多数团队的主要瓶颈将是数据准备和评估,而非训练吞吐量。
具体操作指南
文档中包含一个在包含78张图像的Rider-Waite塔罗牌数据集上对FLUX.1-dev进行端到端微调的完整示例。工作流是:将数据集预编码为缓存的VAE潜在向量和文本嵌入,然后使用现有的FLUX YAML通过命令行覆盖数据集路径和运行设置来启动训练。经过200个优化步骤后,检查点生成的输出将宇航员提示的内容与从塔罗牌中学到的复古色调和墨线轮廓融为一体,这是一种受控的领域适应,而非基础模型接管。
trtcrd触发令牌作为一个风格开关。使用它时,生成的图像会转变为奶油色、红色和黑色的平面色块。不使用它时,相同的种子和提示会产生照片般的效果。这种精细控制正是领域微调应有的样子,而它仅需一个配置文件,无需修改模型源码,这正是关键所在。
不足之处
NeMo Automodel目前仅支持流匹配模型。这排除了社区中仍广泛使用的离散时间扩散模型,例如Stable Diffusion 3和SDXL类架构。该库还配备了SLURM多节点编排,但尚未支持Kubernetes,这限制了它在托管云平台上运行的团队的吸引力。
计划中的Pythonic配方API(允许用户通过类型化Python而非YAML来组合相同组件)仍列为未来版本。对于需要将训练集成到现有实验管理框架中的团队来说,这种程序化路径将比YAML快速启动更重要。
但这并不意味着目前实现的价值打了折扣。Diffusers-NeMo集成解决了一个特定问题:无需常规脚手架进行生产级分布式扩散训练,并且解决得干净利落。性能数字强劲,工作流程可复现,生成的工件与现有工具兼容。对于一直在集群上运行临时脚本的实验室来说,这很可能是通往可扩展设置的最佳路径。
每天早晨用 3 分钟掌握科技要闻
每个工作日一封邮件,只讲真正重要的 AI 与科技动态。