它可以在单个 gpu 上训练超过 130 亿个参数的模型，与 pytorch 等流行框架相比规模增加了 10 倍，而且它不需要数据科学家做任何模型更改，也不需要牺牲计算效率。
资料来源：jie ren 等人在上的文章
巨型模型的训练过程优化正成为 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 上）放在一起的，这很合理。
他们做了几个实验来证明这一点，结果内存使用量有了惊人的缩减！这是这篇论文中我最喜欢的部分之一：他们把问题绘成了一个图，然后分配正确的节点，再剖析问题。
单步延迟参数更新技巧
即使在优化了 cpu 和 gpu 之间的通信后，cpu 的算力限制仍然会带来很大的开销。为了解决这个问题，他们引入了一个新颖的参数延迟技巧。
我们开发了单步延迟参数更新（dpu），将 cpu 和 gpu 的计算重叠起来，通过单步延迟参数更新来隐藏 cpu 的计算开销。我们在评估中验证了 dpu 不会影响训练的最终准确性。
资料来源：jie ren 等人在 arxiv 上的文章
为了进一步优化 cpu 的计算过程，他们提出了一个优化的算法。如果你不熟悉 adam（自适应矩估计）算法的话这里简单说一下：它是机器学习中最经常使用的优化算法之一。他们针对 cpu 优化的 adam 算法需要将[1]参数、梯度、动量和方差放入 cpu 的矢量寄存器。这些寄存器本质上是 cpu 中的存储区域，包含专门为单指令多数据矢量计算设计的操作数。它们也在矢量模式下运行标准的算术运算，并以平铺的方式将所需的数据移动到 gpu 上。这实质上意味着它们是在不重叠的“开放窗口”中移动数据的。
论文中还提到了很多其他方面的优化，但我不想涵盖论文中的每一个细节，我只想给大家展示我眼中最有用和最创新的技巧。
最后的想法和收获
关于巨型模型优化的论文最近频频出新，主要是[1]因为巨型 transformer 的使用率有了惊人的增长。此外，这些 transformer 的大小在不断增加，可数据科学家不应该依赖昂贵的 gpu 和硬件来解决相关问题。我将很快写另一篇文章来对比其中一些巨型模型，因为我目前正在思考这个问题。
参考文献
[1]zero-offload：让十亿规模的模型训练走向大众化。jieren、samyam rajbhari、reza yazdani aminabadi、olatunji ruwasea、、shuangyan yang、minjia zhang、dong li、yuxiong he。2021 发表于
原文链接：
评论