QCon北京「鸿蒙专场」火热来袭!即刻报名,与创新同行~ 了解详情
写点什么

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:3011848
用户头像
刘燕 InfoQ高级技术编辑

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

关注

评论

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

你代码的异味是故意的还是不小心?是故意的!

禅道项目管理

2022 IoTDB Summit:天谋科技刘海《Timecho:孵化于Apache IoTDB的开源商业化公司及其核心产品》

Apache IoTDB

大数据 分布式 开源数据库 IoTDB

探索以小程序提升运维效率

Onegun

运维 小程序容器

等保联盟是什么机构?有什么用?

行云管家

等保 等级保护 等保联盟

Dubbo Triple 协议

昵称不能为null

dubbo RPC triple协议

ByteHouse实时导入技术演进

字节跳动数据平台

大数据 数据仓库 Clickhouse 数据仓库服务 企业号 3 月 PK 榜

性能平台数据提速之路

百度Geek说

大数据 百度 提效 企业号 3 月 PK 榜 性能平台

京东云RASP云原生安全免疫创新实践

京东科技开发者

Web 安全 漏洞 业务安全 企业号 3 月 PK 榜

Higress on K8s 5分钟开箱即用

阿里巴巴中间件

阿里云 云原生 Higress

2022 IoTDB Summit:用友郭关飞《用友在 Apache IoTDB 应用与生态建设方面的探索与实践》

Apache IoTDB

大数据 开源 时序数据库 IoTDB

Vineyard 论文被 SIGMOD'2023 接收,助力计算引擎之间高效数据交换

阿里巴巴中间件

阿里云 计算引擎

SSH和VNC的四个区别简单讲解

行云管家

SSH 行云管家

MMMBSC互助基金系统开发智能合约部署

薇電13242772558

智能合约 dapp

搬得进来,搬得出去!快来过一把数据迁移的“瘾”

OceanBase 数据库

数据库 oceanbase

融云入选中国信通院《高质量数字化转型产品及服务全景图》

融云 RongCloud

产品 数字化 通讯

云智慧助力中国信通院组装式应用开发平台系列标准建设

云智慧AIOps社区

面试官:请说一下如何优化结构体的性能?

王中阳Go

Go golang 高效工作 学习方法 高效学习

技术沙龙 | 探索软件测试前沿技术及最佳实践,体验ChatGPT在测试领域中的应用

测吧(北京)科技有限公司

测试

用户分享 | 达梦第三方客户端DockQuery使用体会

BinTools图尔兹

数据库 用户体验 国产数据库工具

OceanBase 生态产品:时序数据库CeresDB 正式发布 1.0 版本

OceanBase 数据库

数据库 oceanbase

TypeScript 与 JavaScript:你应该知道的区别

京东科技开发者

JavaScript typescript 前端 后端 企业号 3 月 PK 榜

聊聊线上发布这件事

老张

软件测试 权限管理 服务部署

小程序技术如何提升企业的移动研发效率?

FinFish

降本增效 小程序容器 移动研发 小程序技术

动转静两大升级!一键转静成功率领先,重点模型训练提速18%+

飞桨PaddlePaddle

人工智能 百度 飞桨 PaddlePaddle 框架解析

瓴羊Quick BI:“3端4表4擎、3+NPlus”得到众多行业内部人士的认可!

对不起该用户已成仙‖

秒懂算法 | 搜索基础

TiAmo

DFS BFS 搜索算法

React等前端框架如何与小程序结合

Onegun

前端 前端框架 React Vue 3

数据测试实践:从一个bug开始的大数据引擎兼容性探索

京东科技开发者

大数据 bug修复 引擎 测试数据构造 企业号 3 月 PK 榜

精选案例 | 博睿数据30w+监测节点护航新华网、人民网两会重保工作

博睿数据

可观测性 智能运维 博睿数据 精选案例 主动式拨测

NFTScan:怎么使用 NFT API 开发一个 NFT 数据分析平台?

NFT Research

数据分析 API NFT

技术沙龙 | 探索软件测试前沿技术及最佳实践,体验ChatGPT在测试领域中的应用

测试人

软件测试 沙龙 ChatGPT

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