大型神经网络是人工智能领域许多最新进展的核心，但训练它们是一个困难的工程和研究挑战，需要协调一个GPU集群来进行单一的同步计算。随着集群和模型规模的增长，机器学习从业者已经开发出越来越多的技术，在许多GPU上并行化模型训练。乍一看，理解这些并行技术似乎令人生畏，但只要对计算的结构做一些假设，这些技术就会变得更加清晰--在这一点上，你只是把不透明的比特从A地穿梭到B地，就像网络交换机穿梭于数据包一样。
数据并行主义
管道并行化
张量并行化
专家并行主义
数据并行主义
Pipeline Parallelism
张量平行主义
专家并行制
三层模型上各种并行策略的说明。每种颜色指的是一个层，虚线分开不同的GPU。
无并行性
训练一个神经网络是一个迭代的过程。在每一次迭代中，我们通过模型的层向前传递，为一批数据中的每个训练实例计算输出。然后再通过各层向后传递，通过计算每个参数的梯度来传播每个参数对最终输出的影响程度。该批数据的平均梯度、参数和一些每个参数的优化状态被传递给一个优化算法，如Adam，它计算下一个迭代的参数（在你的数据上应该有更好的性能）和新的每个参数优化状态。随着训练在成批数据上的迭代，模型不断进化，产生越来越精确的输出。
各种并行技术在不同的维度上对这个训练过程进行切割，包括。
- 数据并行--在不同的GPU上运行批次的不同子集。
- 管线并行--在不同的GPU上运行模型的不同层。
- 张量并行--将单一操作（如矩阵乘法）的数学运算拆分到不同的GPU上。
- 专家混合--只用每层的一部分来处理每个例子。
(在这篇文章中，我们将假设你使用GPU来训练你的神经网络，但同样的想法适用于那些使用任何其他神经网络加速器的人）。
数据并行化
数据并行训练意味着将相同的参数复制到多个GPU（通常称为 "工作者"），并将不同的例子分配给每个GPU同时进行处理。仅仅是数据并行仍然要求你的模型适合于单个GPU的内存，但让你利用许多GPU的计算，代价是存储许多重复的参数副本。也就是说，有一些策略可以增加你的GPU可用的有效RAM，比如在两次使用之间将参数暂时卸载到CPU内存。
当每个数据并行工作者更新其参数副本时，他们需要协调以确保每个工作者继续拥有类似的参数。最简单的方法是在工作器之间引入阻塞式通信：（1）在每个工作器上独立计算梯度；（2）在各工作器上平均梯度；（3）在每个工作器上独立计算相同的新参数。步骤（2）是一个阻塞平均，需要传输相当多的数据（与工作者的数量乘以参数的大小成正比），这可能会损害你的训练吞吐量。有各种异步同步方案来消除这种开销，但它们会损害学习效率；在实践中，人们通常坚持使用同步方法。
管线并行
通过管道并行训练，我们将模型的顺序块分割到GPU上。每个GPU只持有一部分参数，因此，同一个模型在每个GPU上消耗的内存比例较小。
将一个大型模型分割成连续层的大块是很直接的。然而，各层的输入和输出之间存在着顺序上的依赖性，所以一个天真的实现会导致大量的空闲时间，而工作者在等待前一台机器的输出被用作其输入。这些等待时间被称为 "泡沫"，浪费了闲置机器可以完成的计算。
前进
后退
梯度更新
闲置
一个天真的流水线并行设置的插图，模型按层垂直分割成4个分区。工作者1负责网络第一层的模型参数（最接近输入），而工作者4负责第四层（最接近输出）。"F"、"B"、"U "分别代表前向、后向和更新操作。下标表示一个操作在哪个工作器上运行。由于顺序依赖性，数据一次由一个工作器处理，导致了大量的 "泡沫 "空闲时间。
我们可以重新利用数据并行的思想，通过让每个工作器一次只处理一个数据元素的子集来减少泡沫的成本，使我们能够巧妙地将新的计算与等待时间重叠起来。核心思想是将一个批次分成多个微批；每个微批的处理速度应该是成比例的，每个工作者在下一个微批可用时就开始工作，从而加速管道的执行。有了足够的微批，工人可以在大部分时间内被利用，在步骤的开始和结束时，泡沫最小。梯度是跨微批的平均数，参数的更新只发生在所有微批完成之后。
模型被分割的工作者的数量通常被称为管道深度。
在前向传递期间，工作者只需要将其大块层的输出（称为激活）发送给下一个工作者；在后向传递期间，它只将这些激活的梯度发送给前一个工作者。如何安排这些传递以及如何在微批中聚合梯度，有一个很大的设计空间。GPipe让每个工作者连续处理前向和后向的传递，然后在最后同步聚合来自多个微批的梯度。而PipeDream则安排每个工作者交替地处理前向和后向通道。
前向
后退
更新
闲置
GPipe
管道梦
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等人，2017)
GShard将一个MoE Transformer扩展到6000亿个参数，其方案是只有MoE层被分割到多个TPU设备上，而其他层是完全重复的。Switch Transformer通过将一个输入路由到一个专家，将模型规模扩展到数万亿个参数，甚至更高的稀疏度。
其他节省内存的设计
还有许多其他的计算策略，使训练越来越大的神经网络更容易操作。比如说。
-
为了计算梯度，你需要已经保存了原始激活，这可能会消耗大量的设备内存。 检查点(也被称为激活再计算)存储任何激活的子集，并在后向通道中及时重新计算中间的激活。这可以节省大量的内存，而计算成本最多就是增加一个完整的前向传递。我们也可以通过选择性的激活再计算来不断地在计算和内存成本之间进行权衡，也就是对那些存储成本相对较高但计算成本较低的激活子集进行检查。
-
混合精度训练训练模型是使用较低精度的数字（最常见的是FP16）。现代加速器可以用低精度数字达到更高的FLOP数，而且你还可以节省设备RAM。通过适当的处理，所产生的模型几乎不会失去精度。
-
卸载是将未使用的数据暂时卸载到CPU或不同的设备中，然后在需要时再将其读回。天真的实现方式会大大降低训练速度，但复杂的实现方式会预先获取数据，使设备永远不需要等待数据。这个想法的一个实现是ZeRO，它将参数、梯度和优化器状态分割到所有可用的硬件上，并根据需要将它们具体化。
-
已经提出了内存效率优化器，以减少优化器维护的运行状态的内存占用 ，如Adafactor。
在OpenAI，我们正在训练和改进大型模型，从底层基础设施一直到为现实世界问题部署它们。如果你想把这篇文章中的想法付诸实践--特别是与我们的扩展和应用研究团队有关--我们正在招聘!
鸣谢
感谢Nikolas Tezak、Sam Altman、Daniel Gackle、Ilya Sutskever和Steve Dowling对草案的反馈。感谢Justin Jay Wang, Bianca Martin, 和Steve Dowling的沟通和设计。
