7天花5万美元，我们成功复制了 Stable Diffusion，成本大降88%！训练代码已开源

作者 | Mihir Patel, Cory Stephenson, Landan Seguin

译者 | 核子可乐

策划 | 刘燕

4 月 26 日，AI 创企 Mosaic ML 表示：

我们已经成功用不到 5 万美元复制了 Stable Diffusion，并将训练代码向大家开放！这样的成本水平只相当于我们之前试水项目的三分之一，更是 Stable Diffusion 2 base 本体的八分之一。换言之，每个人都能以前所未有的极低门槛训练出属于自己的 Stable Diffusion。

Mosaic ML 在一篇文章中详细讲述了“复制 Stable Diffusion”的方法、过程以及结果。以下是全文，经 InfoQ 翻译。

“MosaicML 非常适合训练 diffusion 模型，而且相较于以往的工具有了巨大改进。”—— Tony Francis, Dream3D 公司 CEO‘

几个月前，我们曾演示过如何以低廉价格在 MosiacML 平台上从零开始训练大规模 difussion 模型。

今天（4 月 26 日），我们很高兴能为大家带来新的好消息：使用 MosaicML 平台，我们以不到 5 万美元成本花 7.45 天从零开始成功复制了 Stable Diffusion 2。

图一：AI 想象出的菌丝体时装。这种将奇异图像引入设计流程的尝试有望突破创意的边界。以上各图均由我们在 MosaicML 平台上从零训练而成的内部 diffusion 模型创作而成。

利用自有数据训练属于自己的图像生成模型，这个前不久还属痴人说梦的目标如今已经切实可行。通过训练自有 diffusion 模型，我们可以：

1. 使用专有数据；
2. 调整某些艺术或摄影风格的表现形式；
3. 避免违反知识产权法，确保模型能够用于商业用途。我们已经对训练 diffusion 模型所使用的代码和方法进行开源，可供您随意训练自己的模型（https://github.com/mosaicml/diffusion）。

设置

图二：发挥创造力并拥抱意外发现。我们的 diffusion 模型能够生成不同主题、艺术和摄影风格的画面。

模型：我们的 diffusion 模型是一个由变分自动编码器（VAE）、CLIP 模型、U-Net 和扩散噪声调度器组成的 Composer Model，所有功能组件均来自 HuggingFace 的 Diffusers 库。全部模型配置均基于 stabilityai/stable-diffusion-2-base。

图三：Diffusion 模型的简化图示。

数据： 我们使用的训练数据集为 LAION-5B 的一个子集，其中包括带有纯英文标题且审美得分为 4.5+ 的样本。与 Stable Diffusion 2 base 类似，我们根据训练数据的图像分辨率将训练过程划分成两个阶段。在第一阶段，我们使用的是分辨率大于等于 256 x 256 的图像，总计 7.9 亿个图像 - 标题样本。在第二阶段中，我们仅使用分辨率大于等于 512 x 512 的图像，总计 3 亿个图标 - 标题样本。

计算： 两个训练阶段均在 128 个英伟达 A100 GPU 上运行。第一个训练阶段耗时 1.6 天，共运行了 55 万次迭代；第二阶段耗时 4.9 天，共运行了 85 万次迭代，总训练时长为 20051 个 A100 小时。除了训练时间之外，我们还预先计算了 VAE 和 CLIP 模型的潜伏空间，希望减少数据集多次传递所带来的训练时间和成本。潜伏空间计算大致需要额外 3784 个 A100 小时，所以模型的总训练时长为 23835 个 A100 小时。假设 A100 的使用成本为每小时 2 美元，则总价格为 4.77 万美元。

技术栈： 我们使用 Composer 作为训练框架，使用 StreamingDataset 来加载 100 TB 训练数据，并使用 MosaicML 平台解决 128 个 GPU 作为训练和评估基础设施时的部署和管理挑战。

图四：训练过程中的损失曲线。我们的平台发现了两个硬件故障，并在无人干预的情况下自动完成了重启。损失曲线之所以不连续，是因为第二阶段将分辨率从上阶段的 256 x 256 提高到了 512 x 512。

挑战与解决方案

无论是 diffusion 扩散模型还是大语言模型，规模化训练都需要经历一系列重大挑战。我们使用 MosaicML 平台进行 diffusion 模型训练，该平台自动解决了大部分问题，确保我们能专注于训练出最佳模型。下面是规模化训练中的三个主要挑战，还有我们的平台如何加以解决。

基础设施

在大规模数据集上训练大模型无疑需要海量算力。MosaicML 平台能够轻松在任意云服务商处编排数百个 GPU。例如，我们的主训练作业运行在一个包含 128 个 A100 GPU 的集群当中。为了确保评估模型不会拖慢训练速度，我们使用不同云服务商在不同集群的各个检查点上自动启用运行评估，并根据可用性将运行规模收缩至 64 乃至最少 8 个 GPU 上。

即使是在训练开始之后，软件或硬件故障也有可能导致训练中断，这就要求 24/7 全天候加以监控。好在 MosaicML 平台的 Node Doctor 和 Watchdog 功能会自动检测故障节点，并根据需要执行恢复操作。通过自动恢复，我们得以从故障中顺利恢复，无需任何人为干预即可继续训练，避免了昂贵的停机时间和人工管理。启动之后，一切无忧！

软件效率

软件配置的优化向来是个大麻烦，好在我们基于 PyTorch 的 Composer 库能够最大程度提高训练效率。跟上一轮实验类似，随着 GPU 数量的增加，Composer 继续保持着出色的吞吐量扩展能力。在本次更新中，我们添加了进一步优化（低精度 GroupNorm 和低精度 LayerNorm，全分片化数据并行）以实现近乎完美的强大扩展能力，将作业最多扩展至 128 个 GPU，从而将成本控制在 5 万美元以内。我们还使用 Composer 的原生指数移动平均（EMA）算法，得以在接近训练结束时（第二阶段的 80 万次迭代中）启用 EMA，从而节约下相当一部分内存和训练算力。

管理 100 TB 数据

我们在训练中使用的是包含 7.9 亿个样本的 LAION-5B 子集，总数据量超过 100 TB。庞大的数据集规模导致其难以管理，特别是在需要配合拥有独立本地存储的多集群情况下。

MosaicML StreamingDataset 库让海量数据集的处理变得更加简单快速，该库提供的三个核心功能也在本次训练中发挥了关键作用：

1. 将存储在不同位置的数据集混合起来。 我们根据图像分辨率将各样本分别存储在不同的数据集内。在训练时，我们使用 MosaicML StreamingDataset 库将来自各数据集的分辨率素材混合起来。

2.即时轮中恢复。 我们能够在一个轮次期间即时恢复训练，这相当于实现了整个数据集在训练过程中的“断点续传”，大大节约了总体用时。

3.以弹性方式实现确定性。MosaicML StreamingDataset 库能够以确定性方式混洗数据，且不受训练用 GPU 数量变化的影响。这使我们得以准确重现训练效果，极大简化了调试步骤。

人类评估结果

图像生成模型的实际性能往往难以评估，除了投入人力别无他法。在盲测评估中，我们衡量了用户对图像质量的偏好，并在 Stable Diffusion 2 和我们自己的 diffusion 模型间进行了提示词对齐。根据用户偏好，我们得出的结论是两套模型质量相当（参见图五）。所有图像均根据 Imagen 论文中提出的 Drawbench 基准测试揭示词生成。

图五：我们对图像质量（左）和揭示词对齐（右）的人工评估结果。误差条显示的置信区间为 95%。在这两次实验中，两套模型之间的用户偏好差异与衡量的不确定性相当，因此我们得出结论认为两套模型的整体质量也基本相当。

未来展望

本文向大家介绍了我们这套 diffusion 模型的输出性能和损失曲线，描述了高级模型训练中的种种细节，还有 MosaicML 平台帮助我们解决的规模化训练挑战。但很遗憾，由于 LAION-5B 数据集使用要求和相关法律条款较为模糊，我们暂时还无法对外公布由此训练出的图像生成模型的参数权重。我们很清楚参数权重对于图像生成模型性能的重要意义，但这里只能向大家说声抱歉。

写在最后

下面就是我们这套 diffusion 模型生成的图像结果。团队成员们都玩得不亦乐乎，也希望各位能从中找到属于自己的乐趣。

原文链接：

https://www.mosaicml.com/blog/training-stable-diffusion-from-scratch-part-2