11

InfoQ · 5月25日

微软提出训练巨型模型的新模式:ZeRO-Offload可训练高达700亿参数的模型

它可以在单个 GPU 上训练超过 130 亿个参数的模型,与 PyTorch 等流行框架相比规模增加了 10 倍,而且它不需要数据科学家做任何模型更改,也不需要牺牲计算效率。资料来源:Jie Ren 等人在 arxiv 上的文章

巨型模型的训练过程优化正成为 ML 领域的一大趋势。

我们已经见识了谷歌的 Switch Transformers 和 OpenAI 的 GPT-3,现在是时候看看微软新出的 ZeRO-Offload 了。该模型 [1] 的设计重点在于通过新的优化技巧,将从 GPU 到 CPU 的卸载成本降到最低。它能够训练参数多达 700 亿的模型,这样的数字实在太令人惊叹了。

这个模型的主要卖点之一是:它经过充分优化,可以在单个 GPU 上进行训练,因此,实际上每个人都可以使用它(你用不着购买大量昂贵的硬件才能使用这个模型)。他们主要通过 异构深度学习 方法来做到这一点,这种方法显著优化了 CPU 与 GPU 之间的协作效率。

背景知识

在我们开始深入研究这个模型之前,我们首先需要了解它所要解决的问题的性质。ML 模型往往会消耗大量内存。为什么?因为它们的参数、梯度和优化器的状态 [1] 在训练期间必须存储在内存中才行。一个常见的解决方案是将这些模型状态分割在多个 GPU 上,然而这种做法的实用性不高,因为大多数数据科学家并没有那么多 GPU 可用。

其他方法包括将状态保存到硬盘上并重用它,或者以某种方式压缩它们。我们将在本文中讨论的方法则是使用 CPU 的内存。这可能算不上是一种新方法,然而这种方法到目前为止并不是很成功,因为在 GPU 和 CPU 之间有很大的通信成本。

这篇论文的一个独特想法体现在他们寻找解决方案的路径上。他们决定使用 数据流图和第一原理分析 来建模这个通信问题。而这方面大部分的工作都是为了在 CPU 和 GPU 设备之间高效地划分这些图。

深入研究模型细节

此外,我们的第一原理分析表明,ZeRO-Offload 在最大限度地节省内存,同时最大限度地减少通信开销和 CPU 计算开销方面为大型模型训练提供了一个唯一的最佳解决方案。ZeRO-Offload 是围绕三大支柱设计的:i)效率,ii)可扩展性,和 iii)可用性。

资料来源:Jie Ren 等人在 arxiv 上的文章

图正在成为机器学习中的一个热门话题。它们也是复杂计算机科学问题解决方案的基本组成部分之一(如果读者感兴趣,我建议看看 Dijkstra 的 算法)。这篇论文使用的图模型将节点之间的数据流表示为边,将数据量表示为这些边的权重。

这里的一个要点是,论文的作者必须充分利用 GPU 和 CPU 各自的主要优势。虽然 CPU 的计算能力通常不是很高,但如果你把合适的计算需求卸载到它上面,你的训练过程就会变得更顺畅。而这正是他们所做的事情:他们将前向和后向传播过程移到了 GPU 上,并将剩余的计算过程卸载到了 CPU 上。他们称前向和后向传播节点为“超级节点”。能看到计算机系统方面的技巧被用来优化机器学习是很令人欣慰的(我是操作系统和 Linux 内核的忠实粉丝!)。这篇论文可以提醒你,机器学习不仅仅是关于数学和数据分析的学科。

下一部分是最小化 CPU 和 GPU 之间的通信量。要做到这一点,你必须考虑如何才能以最优化的方式切割这些图,因为如果你不能正确地分割它,就只能继续向前和向后发送数据。他们提出了一个解决方案,其中的图有 4 个节点:一个超级节点(传播),一个参数数据节点(p16),一个梯度数据节点(g16),以及一个更新超级节点的节点。他们假设超级节点肯定是与参数数据节点(在 GPU 上)放在一起的,这很合理。

1.png

他们做了几个实验来证明这一点,结果内存使用量有了惊人的缩减!这是这篇论文中我最喜欢的部分之一:他们把问题绘成了一个图,然后分配正确的节点,再剖析问题。

单步延迟参数更新技巧

即使在优化了 CPU 和 GPU 之间的通信后,CPU 的算力限制仍然会带来很大的开销。为了解决这个问题,他们引入了一个新颖的参数延迟技巧。

我们开发了单步延迟参数更新(DPU),将 CPU 和 GPU 的计算重叠起来,通过单步延迟参数更新来隐藏 CPU 的计算开销。我们在评估中验证了 DPU 不会影响训练的最终准确性。

资料来源:Jie Ren 等人在 arxiv 上的文章

为了进一步优化 CPU 的计算过程,他们提出了一个优化的 Adam 算法。如果你不熟悉 Adam(自适应矩估计)算法的话这里简单说一下:它是机器学习中最经常使用的优化算法之一。他们针对 CPU 优化的 Adam 算法需要将 [1] 参数、梯度、动量和方差放入 CPU 的 矢量寄存器。这些寄存器本质上是 CPU 中的存储区域,包含专门为单指令多数据矢量计算设计的操作数。它们也在矢量模式下运行标准的算术运算,并以平铺的方式将所需的数据移动到 GPU 上。这实质上意味着它们是在不重叠的“开放窗口”中移动数据的。

论文中还提到了很多其他方面的优化,但我不想涵盖论文中的每一个细节,我只想给大家展示我眼中最有用和最创新的技巧。

最后的想法和收获

关于巨型模型优化的论文最近频频出新,主要是 [1] 因为巨型 transformer 的使用率有了惊人的增长。此外,这些 transformer 的大小在不断增加,可数据科学家不应该依赖昂贵的 GPU 和硬件来解决相关问题。我将很快写另一篇文章来对比其中一些巨型模型,因为我目前正在思考这个问题。

本文转自 公众号:AI前线 ,作者Mostafa Ibrahim,点击阅读原文
11 阅读 298
推荐阅读
关注数
12452
内容数
140
跟踪科技公司、科研机构和IT产业在AI方向的最新动态,挖掘AI技术应用场景和商业化落地案例。
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息