限时领|《AI 百问百答》专栏课+实体书(包邮)! 了解详情
写点什么

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

  • 2023-05-03
    北京
  • 本文字数:2995 字

    阅读完需:约 10 分钟

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


2023-05-03 13:3012061
用户头像
刘燕 InfoQ高级技术编辑

发布了 1112 篇内容, 共 572.4 次阅读, 收获喜欢 1980 次。

关注

评论

发布
暂无评论
发现更多内容

「数」论|厂级实时监控系统SIS:盛宴还是剩宴?

麦杰研究院

一文掌握 1688 商品详情 API 接口:从入门到实战

tbapi

1688商品详情接口 1688API 1688数据接口

AI智上 | 数智报告:500强央企集团合并报表领先实践

用友智能财务

用友 数智化 智能财务 合并报表

IoTDB 广州行 | 4.26 端边云协同的新一代 AI 数据库技术沙龙回顾

Apache IoTDB

不用Mockito写单元测试?你可能在浪费一半时间

Java随想录

Java 单元测试 Mockito

YashanDB|Windows 下如何配置 ODBC 跟踪日志?三步搞定!

数据库砖家

数据库

A2A 是否只是换个了名字的 MCP?开发者该如何选择?

Baihai IDP

程序员 AI agent MCP A2A

机械加工行业QMS质量管理系统

万界星空科技

万界星空科技 QMS 生产质量管理 机械加工行业 产品质量管理

Netflix系统架构解析

码语者

再这样玩下去,以太坊就是离心机,就是吸血魔

ChainFlash链訊

区块链 以太坊

AI与无人零售:如何通过智能化技术提升消费者体验和运营效率?

天津汇柏科技有限公司

人工智能 AI

YashanDB|decode 子查询被多次执行?性能骤降的真相找到了!

数据库砖家

数据库

「数」论|厂级实时监控系统SIS:盛宴还是剩宴?

麦杰科技

AI 绘画党AIGC创作者看过来!AI提示词可以统一管理了

龙正哲

独立开发者工具 • 半月刊 第 008 期

Immerse

SEO与国际化

溪抱鱼

SEO

Qwen3初测,小尺寸MOE利好小玩家

冯骐

人工智能 大模型 大语言模型 AI 基础设施 Qwen3

1688 商品列表 API 深度拆解:从参数配置到数据获取

tbapi

1688商品列表接口 1688商品数据接口 1688API

报名丨第七届智源大会,6月6-7日北京见

智源研究院

怎么做好质量管理?质量管理怎么做到高效又保质?

积木链小链

质量管理 数字化转型 智能制造

PIRF:432 - Performance - Pressure, Precision, Payoff

Echo!!!

English

Qwen3 系列全家桶,百度百舸一键部署

Baidu AICLOUD

AI 基础设施

利用通义灵码和魔搭 Notebook 环境快速搭建一个 AIGC 应用 | 视频课

阿里巴巴云原生

阿里云 云原生 通义灵码

StarRocks 查询优化器深度解析

StarRocks

浅谈如何使用 Amazon CUR 报告做更完善的成本分析

亚马逊云科技 (Amazon Web Services)

鸿蒙游戏生态加速:华为小游戏焕新升级,抢滩百亿级市场新机遇

最新动态

英特尔代工:明确重点广合作,服务客户铸信任

E科讯

黑龙江等保测评应急预案

黑龙江陆陆信息测评部

啡漫AIP首发!AI+咖啡+动漫跨界融合,解锁咖啡产业AI新玩法

科技热闻

JMeter、Apipost 与 Postman 的 API 测试对比:为什么 APIPost 是更聪明的选择

数据追梦人

中国制造的AI穿凿,一场直抵地心的技术远征

脑极体

AI

7天花5万美元,我们成功复制了 Stable Diffusion,成本大降88%!训练代码已开源_AI&大模型_Mihir Patel_InfoQ精选文章