华裔女工程师Lilian Weng和CTO Greg Brockman合撰的文章。介绍了在众多GPU上并行训练大模型的各种技术。全文编译如下：

大型神经网络是人工智能许多最新进展的核心，但训练它们是一项艰巨的工程和研究挑战，需要编排一个GPU集群来执行单个同步计算。随着集群和模型规模的增长，机器学习从业者已经开发出越来越多的技术来并行化许多 GPU 上的模型训练。乍一看，理解这些并行技术似乎令人生畏，但是只需对计算结构进行一些假设，这些技术就会变得更加清晰 - 在这一点上，只是像网络交换机在数据包周围穿梭一样，从A到B的比特。

三层模型上各种并行性策略的图示。每种颜色都表示一个图层，虚线分隔不同的 GPU。

 

无并行性

训练神经网络是一个迭代过程。在每次迭代中，我们都会在模型的层中向前传递，以计算一批数据中每个训练示例的输出。然后，另一个过程通过层向后进行，通过计算相对于每个参数的梯度来传播每个参数对最终输出的影响程度。批处理、参数和一些每个参数优化状态的平均梯度将传递给优化算法，例如 Adam，该算法计算下一次迭代的参数（这些数据的性能应该稍好一些）和新的每个参数优化状态。随着训练对批量数据的迭代，模型会不断发展以产生越来越准确的输出。

各种并行技术将此训练过程划分为不同的维度，包括：

  • 数据并行性 - 在不同的 GPU 上运行批处理的不同子集;
  • 管道并行性 — 在不同的 GPU 上运行模型的不同层;
  • 张量并行性 - 分解单个操作的数学，例如要在GPU之间拆分的矩阵乘法;
  • 专家混合 - 仅按每层的一小部分处理每个示例。

（在这篇文章中，我们假设你正在使用GPU来训练你的神经网络，但同样的想法也适用于那些使用任何其他神经网络加速器的人。）

 

数据并行性

数据并行训练意味着将相同的参数复制到多个 GPU（通常称为“worker”），并为每个 GPU 分配不同的示例以进行同时处理。仅数据并行性仍然要求模型适合单个 GPU 的内存，但允许您利用许多 GPU 的计算，但代价是存储许多参数的重复副本。话虽如此，有一些策略可以增加GPU可用的有效RAM，例如在两次使用之间临时将参数卸载到CPU内存。

当每个数据并行工作线程更新其参数副本时，他们需要协调以确保每个工作线程继续具有相似的参数。最简单的方法是引入工作单元之间的阻塞通信：（1）独立计算每个工作单元的梯度;（2）平均工作单元之间的梯度;（3）独立计算每个工作线程上相同的新参数。步骤（2）是一个阻塞平均值，需要传输相当多的数据（与工作线程数乘以参数大小成比例），这可能会损害训练吞吐量。有各种异步同步方案可以消除这种开销，但它们会损害学习效率；在实践中，人们通常坚持使用同步方法。

 

流水线并行

通过管道并行训练，我们将模型的顺序块分区到 GPU 中。每个 GPU 只包含一小部分参数，因此同一模型每个 GPU 消耗的内存比例更低。

将大型模型拆分为连续层块非常简单。但是，层的输入和输出之间存在顺序依赖关系，因此朴素的实现可能会导致大量的空闲时间，而工作线程等待前一台计算机的输出用作其输入。这些等待时间块被称为“气泡”，浪费了空转机器可以完成的计算。

简单管道并行度设置的图示，其中模型按层垂直拆分为 4 个分区。工作线程 1 承载网络第一层（最接近输入）的模型参数，而工作线程 4 承载第 4 层（最接近输出）。“F”、“B”和“U”分别表示前进、后退和更新操作。下标指示操作在哪个工作线程上运行。由于顺序依赖性，数据一次由一个工作线程处理，从而导致空闲时间的大“气泡”。

我们可以重用数据并行性中的想法，通过让每个工作线程一次只处理一个数据元素的子集来降低泡沫的成本，从而允许我们巧妙地将新计算与等待时间重叠。核心思想是将一个批次分成多个微粒;每个微粒的处理速度应该成比例地快，每个工作单元在下一个微粒可用时立即开始工作，从而加快管道执行。通过足够的微粒，工作单元可以在大部分时间里使用，在步骤的开始和结束时有最小的气泡。梯度在微微熔池之间取平均值，并且只有在所有微熔池完成后才会更新参数。

模型被分割的工作单元数通常称为管道深度。

在正向传递期间，工作线程只需将其层块的输出（称为激活）发送到下一个工作线程;在向后传递期间，它仅将这些激活的梯度发送到上一个工作线程。关于如何安排这些通道以及如何跨微切块聚合梯度，有很大的设计空间。GPipe让每个工作单元进程向前和向后连续通过，然后在最后同步聚合来自多个微切口的梯度。相反，PipeDream 会安排每个工作线程交替地处理向前和向后传递。

GPipe和PipeDream流水线方案的比较，每批使用4个微粒。微切 1-8 对应于两个连续的数据批次。在图像中，“（数字）”表示在哪个微批处理上执行操作，下标标记了工作线程 ID。请注意，PipeDream 通过使用陈旧的参数执行一些计算来提高效率。

 

张量并行性

管道并行性按层“垂直”拆分模型。也可以“水平”拆分层内的某些操作，这通常称为张量并行训练。对于许多现代模型（如 Transformer），计算瓶颈是将激活批处理矩阵与大权重矩阵相乘。矩阵乘法可以被认为是行和列对之间的点积;可以在不同的 GPU 上计算独立的点积，或者在不同的 GPU 上计算每个点积的部分并总结结果。无论采用哪种策略，我们都可以将权重矩阵切割成均匀大小的“分片”，将每个分片托管在不同的GPU上，并使用该分片计算整个矩阵产品的相关部分，然后再进行通信以合并结果。

一个例子是Megatron-LM，它在Transformer的自注意层和MLP层内并行化矩阵乘法。PTD-P使用张量，数据和管道并行性；它的管道计划为每个设备分配多个非连续层，以减少更多网络通信为代价的气泡开销。

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

 

专家混合（MOE）

使用专家混合（MoE）方法，只有一小部分网络用于计算任何一个输入的输出。一个示例方法是具有许多组权重，并且网络可以在推理时通过门控机制选择要使用的集合。这样可以在不增加计算成本的情况下实现更多参数。每组权重被称为“专家”，希望网络将学会为每个专家分配专门的计算和技能。不同的专家可以托管在不同的 GPU 上，从而提供了一种扩展用于模型的 GPU 数量的明确方法。

 

专家混合 （MoE） 层的图示。门控网络仅选择 n 个专家中的 2 个。（图片改编自：Shazeer et al.， 2017）

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

 

其他节省内存的设计

还有许多其他计算策略可以使训练越来越大的神经网络更容易处理。例如：要计算梯度，您需要保存原始激活，这可能会消耗大量设备RAM。

  • 检查点（也称为激活重新计算）存储激活的任何子集，并在向后传递期间实时重新计算中间激活的子集。这节省了大量内存，而计算成本最多为一次额外的完全向前传递。人们还可以通过选择性激活重新计算在计算和内存成本之间不断权衡，这是对激活的检查点子集进行检查点操作，这些子集的存储成本相对较高，但计算成本较低。
  • 混合精度训练是使用较低精度的数字（最常见的是FP16）训练模型。现代加速器可以用较低的精度数字达到更高的 FLOP 计数，并且您还可以节省设备 RAM。如果小心谨慎，生成的模型几乎不会丢失任何准确性。
  • 卸载是暂时将未使用的数据卸载到CPU或不同设备之间，然后在需要时将其读回。朴素的实现会大大减慢训练速度，但复杂的实现会预取数据，以便设备永远不需要等待它。这个想法的一个实现是ZeRO，它将参数，梯度和优化器状态拆分到所有可用的硬件中，并根据需要实现它们。
  • 内存高效优化器来减少由优化器（如Adafactor）维护的运行状态的内存占用。
  • 压缩还可用于在网络中存储中间结果。例如，Gist 压缩为向后传递保存的激活；DALL·E 在同步渐变之前对其进行压缩。

内容中包含的图片若涉及版权问题，请及时与我们联系删除