AICon上海|与字节、阿里、腾讯等企业共同探索Agent 时代的落地应用 了解详情
写点什么

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

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

关注

评论

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

宜搭5月更新:跨应用数据读写能力升级,AI组件内测开放

一只大光圈

美团二面:为什么Redis会有哨兵?

Java全栈架构师

Java 数据库 redis 程序员 面试

源码解读预告 |TiFlash DeltaTree 引擎设计及实现解析!

TiDB 社区干货传送门

【LeetCode】数组中的第K个最大元素Java题解

Albert

LeetCode 5月月更

4月券商App行情刷新及交易体验评测报告,7家券商入围领导者象限

博睿数据

性能测试 系统运维 博睿数据 券商排行

SysAK 应用抖动诊断篇—— eBPF又立功了! | 龙蜥技术

OpenAnolis小助手

Linux 工具 内核 ebpf 龙蜥技术

明天,龙蜥2位专家直播,第22届计算机系统会议等活动来了!

OpenAnolis小助手

Linux 开源 直播 内核 龙蜥技术

linux之awk使用技巧

入门小站

css基本概念学习篇【四】

恒山其若陋兮

5月月更

springboot集成activiti整套方案()

金陵老街

Vue ERP Activiti spring-boot

天翼云推荐新人返好礼,最高返利千元

天翼云开发者社区

Redis「9」主从、高可用性方案

Samson

redis 学习笔记 5月月更

AIRIOT物联网低代码平台如何配置OPC UA驱动?

AIRIOT

Authing 被世界经济论坛评选为 2022 技术先锋企业

Authing

身份云 科技 Idaas 科技企业

生命科学领域下的医药研发通过什么技术?冷冻电镜?分子模拟?IND?

GPU算力

数据库连接池 -Druid 源码学习(十)

wjchenge

Druid 数据库连接池

大家谈的视频体验指标,都有哪些?如何测定?

声网

视频 Qoe Dev for Dev

ECSM隐私协议

潇潇雨歇

All in ONE!博睿数据重磅推出一体化智能可观测平台

博睿数据

博睿数据 IT运维 ONE平台

半年面试数百场,我总结出了这份10w字Java面试复盘笔记

Java全栈架构师

Java spring 程序员 架构 面试

windows下C语言使用curl库访问HTTP下载文件

DS小龙哥

5月月更

5.26直播预告|《观见话题》第一期:跨境组网与加速上云的硬核解法

观测云

极速调取客户保单,YRCloudFile 助力保险存储架构升级

焱融科技

AI 存储 NAS 数字金融

普渡科技宣布成立“PUDU-X”创新基金,为青年工程师筑梦未来

极客天地

【活动报名】TiDB 社区天津站 Meetup 要来啦!

TiDB 社区干货传送门

青藤入选信通院“数据安全推进计划”成员单位

青藤云安全

数据安全 信通院

JSON在线对比差异工具

入门小站

工具

OceanBase 源码解读(十一):Location Cache 模块浅析

OceanBase 数据库

oceanbase 源码解读

【高并发】什么是ForkJoin?看这一篇就够了!

冰河

并发编程 多线程 高并发 协程 异步编程

Authing 身份云招聘:增长黑客

Authing

招聘 科技

使用 Provider 实现 Flutter 多组件的状态共享

岛上码农

flutter 安卓开发 ios 开发 跨平台应用 5月月更

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