巨型模型的训练过程优化正成为 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算法。如果你不熟悉 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 发表于arxiv
原文链接：
