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

近日 | OpenAI

编译 | 黄楠

感谢 | 陈彩娴

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

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

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

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

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

可能混合（MOE）——仅通过每层得一小部分处理每个示例。

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

1 数据并行

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

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

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

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

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

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

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

2 流水线并行

在流水线并行训练中，研究者会将模型得顺序块划分到 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 发布者会员账号。注意，PipeDream 通过使用陈旧参数执行一些计算来获得更高得效率。

3 张量并行

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

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

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

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

4 可能混合 (MoE)

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

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

图注：门控网络只选择了n个可能中得2个。

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

5 其他节省内存得设计

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

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

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

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

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

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

参考链接：

感谢分享openai感谢原创分享者/blog/techniques-for-training-large-neural-networks/