OpenAI：训练大型神经网络的四种基本方法

来源 | OpenAI

编译 | 黄楠

大型神经网络是当前人工智能领域的热门话题之一，那么，如何训练大模型？

最近，曾推出大规模预训练模型 GPT-3 的 OpenAI 发表了一篇博文，介绍了基于 GPU 的四种节省内存的并行训练方法，分别是：

• 数据并行——在不同的 GPU 上运行同一批次的不同子集；

• 流水线并行——在不同的 GPU 上运行模型的不同层；

• 张量并行——分解单个运算的数学运算，例如将矩阵乘法拆分到 GPU 上；

• 专家混合（MOE）——仅通过每层的一小部分处理每个示例。

图注：三层模型上各种并行策略，每种颜色代表一层，虚线分隔不同的 GPU。

「数据并行训练」意味着将相同的参数复制到多个 GPU（通常称为“workers”），并为每个 GPU 分配不同的示例以同时处理。

单单的数据并行要求模型匹配单个 GPU 内存，但当你利用多个 GPU 计算时，代价是存储参数的多个副本。不过，话虽如此，有一些策略可以增加 GPU 可用的有效 RAM，例如，在两次使用之间，可将参数暂时卸载到 CPU 内存。

随着每次数据并行 worker 更新其参数副本，它们需要相互协调，以确保每个 worker 都继续具有相似的参数。最简单的方法是在 worker 之间引入「阻塞通信」：

步骤 1：独立计算每个worker上的梯度；

步骤 2：将不同 worker 的梯度平均；

步骤 3：在每个 worker 上独立计算相同的新参数。

步骤 2 是一个阻塞平均值，它需要传输大量数据（与 worker 数量乘以参数大小成正比），这可能会损害训练的吞吐量。有各种异步同步方案可以消除这种损耗，但会损害学习效率；因此在实践中，人们普遍坚持同步方法。

在流水线并行训练中，研究者会将模型的顺序块划分到 GPU 上，每个 GPU 只保存一小部分参数，因此，相同模型的每个 GPU 消耗的内存按比例减少。

将大型模型拆分为连续层的块很简单，但由于层的输入和输出之间存在顺序依赖关系，因此，在 worker 等待前一台机器的输出用作其输入时，一个幼稚的执行可能会导致出现大量空闲时间。这些等待时间块被称为「泡沫」（bubbles），即浪费了本可以由空闲机器来完成的计算。

图注：一个简单的流水线并行设置插图，其中，模型被垂直分成 4 个分区。worker 1 主持第一层的模型参数（最接近输入），而 worker 4 主持第 4 层（最接近输出）。“F”、“B”和“U”分别代表前向、后向和更新操作。下标会指示在哪个 worker 上运行操作。由于顺序依赖性，数据一次由一个 worker 处理，导致产生了大量的空闲时间“泡沫”。

我们可以重用数据并行的想法，通过让每个 worker 一次只处理数据元素的一个子集，来降低产生时间泡沫的成本，从而使我们能巧妙地将新计算与等待时间重叠。核心思想是，将一个批次拆分为多个微批次，每个微批次的处理速度都应该成比例地加快，并且每个 worker 在下一个微批次可用时立即开始工作，从而加快管道执行。有了足够的微批次， worker 可以在大部分时间被利用，并且在步骤开始和结束时「泡沫」最小。梯度在微批次之间进行平均，并且只有在所有微批次完成后才会更新参数。

模型拆分的 worker 数量通常称为「管道深度」（pipeline depth）。

在前向传递期间，worker 只需将其层块的输出（称为「激活」）发送给下一个 worker；在反向传递期间，它仅将这些激活的梯度发送给前一个工作人员。如何安排这些通道以及如何跨微批次聚合梯度有很大的设计空间。例如，方法 GPipe 是让每个工作进程连续向前和向后传递，然后在最后同步聚合来自多个微批次的梯度；而 PipeDream 会安排每个 worker 交替处理的前向和后向通道。

图注：GPipe 和 PipeDream 流水线方案的比较，每批使用 4 个微批次。微批次 1-8 对应于两个连续的数据批次。图中“number”表示在哪个微批次上操作，下标标记 worker  ID。注意，PipeDream 通过使用陈旧参数执行一些计算来获得更高的效率。

管道并行性将模型逐层“垂直”拆分，也可以在一个层内“水平”拆分某些操作，这通常称为张量训练。

对于许多现代模型（例如Transformer），计算瓶颈是将激活批处理矩阵与大权重矩阵相乘。矩阵乘法可以认为是成对的行和列之间的点积；可以在不同的 GPU 上计算独立的点积，或者在不同的 GPU 上计算每个点积的部分并总结结果。无论采用哪种策略，我们都可以将权重矩阵分割成大小均匀的“碎片”，将每个碎片托管在不同的 GPU 上，并使用该碎片计算整个矩阵乘积的相关部分，然后再进行通信以组合结果。

一个例子是Megatron-LM，它在 Transformer 的自注意力和 MLP 层内并行化矩阵乘法。PTD-P使用张量、数据和流水线并行，其流水线调度为每个设备分配了多个不连续的层，以增加网络通信为代价来减少泡沫损耗。

有时，网络输入可以跨维度并行化，相对于交叉通信具有高度的并行计算。序列并行就是这样一种想法，其中输入序列在时间上被分成多个子示例，通过允许计算继续进行更细粒度的示例，来按比例减少峰值内存消耗。

使用专家混合（MoE）方法，只有小部分网络用于计算任何一个输入的输出。

一个示例方法是拥有多组权重，并且网络可在推理时通过门控机制选择要使用的权重组，这能在不增加计算成本的情况下启用更多参数。每组权重都被称为“专家”，且希望网络能学会为每个专家分配专门的计算和技能。不同的专家可以主持不同的 GPU ，从而提供了一种明确的方式来扩大用于模型的 GPU 数量。

图注：门控网络只选择了n个专家中的2个。

GShard 将 MoE Transformer 的参数扩展到 6000 亿个参数，其中仅将 MoE 层拆分到多个 TPU 设备上，其他层则完全复制。Switch Transformer 通过将一个输入路由到单个专家，将模型大小扩展到数万亿个参数，具有更高的稀疏性。

还有许多其他的计算策略，可以使训练越来越大的神经网络更容易处理。例如：

要计算梯度，需要保存原始激活，这会消耗大量设备 RAM。检查点（也称为激活重新计算）存储激活的任何子集，并在反向传递期间，及时重新计算中间的激活，以最多一个额外完整前向传递的计算成本，节省了大量内存。人们还可以通过选择性激活重新计算，来不断权衡计算和内存成本，这是对激活的子集进行检查，其存储成本相对较高，但计算成本较低。

混合精度训练是使用较低精度的数字（最常见的是FP16）来训练模型。现代加速器可以使用较低精度的数字达到更高的 FLOP 计数，并且还能节省设备 RAM。在适当的照顾下，产生的模型几乎可以不损失任何精度。

卸载是将未使用的数据临时卸载到 CPU 或不同设备之间，在需要时将其读回。幼稚的执行会大大减慢训练速度，但复杂的实现方式会预先获取数据，使设备永远不需要等待。这个想法的一个实现是ZeRO，它可将参数、梯度和优化器状态分割到所有可用的硬件上，并根据需要将它们具体化。

Memory Efficient Optimizers已经提出了内存效率优化器，以减少优化器所维护的运行状态的内存占用，例如Adafactor。

压缩也可用于存储网络中的中间结果。例如，Gist压缩为后向传递而保存的激活；DALL-E在同步梯度之前压缩梯度。

参考链接：

雷峰网(公众号：雷峰网)