本文转载自雷锋网。如需转载，请在雷锋网官网申请授权。大型神经网络是目前人工智能领域的热门话题之一，那么如何训练大型模型呢？近日，推出大规模预训练模型GPT-3的OpenAI发表了一篇博文，介绍了四种基于GPU的内存节省并行训练方法，即：数据并行——在不同的GPU上运行同一个batch不同的子集子集；管道并行性——在不同的GPU上运行模型的不同层；张量并行——分解单个操作的数学运算，例如将矩阵乘法拆分到GPU上；混合专家(MOE)-仅通过每一层的一小部分处理每个示例。图注：三层模型上的各种并行策略，每种颜色代表一层，虚线分隔不同的GPU。1数据并行“数据并行训练”是指将相同的参数复制到多个GPU（通常称为“worker”），并为每个GPU分配不同的示例以同时处理。仅数据并行性就需要模型适合单个GPU内存，但是当您利用多个GPU计算时，成本是存储参数的多个副本。然而，话虽如此，还是有一些策略可以增加GPU可用的有效RAM，例如在使用之间临时将参数卸载到CPU内存。当每个数据并行worker更新其参数副本时，它们需要相互协调以确保每个worker继续具有相似的参数。最简单的方法是在worker之间引入“阻塞通信”：步骤1：独立计算每个worker上的梯度；第二步：平均不同worker的梯度；第3步：在每个worker独立参数上计算相同的新梯度。第2步是一个阻塞平均，需要传输大量数据（与工人数量乘以参数大小成正比），这会影响训练吞吐量。有多种异步-同步方案可以消除这种损失，但以牺牲学习效率为代价；因此在实践中，人们通常坚持同步方法。2PipelineParallelism在流水线并行训练中，研究人员将模型的顺序块划分到GPU上，每个GPU只保存一小组参数，因此，每个GPU对同一模型消耗的内存按比例减少。将大型模型拆分为连续层的块很简单，但由于层的输入和输出之间的顺序依赖性，当工作人员等待前一台机器的输出用作其输入时，简单的执行可能会失败。导致大量空闲时间。这些等待时间块被称为“泡沫”，浪费了本可以由闲置机器完成的计算。图例：简单管道并行设置的图示，其中模型垂直分为4个分区。worker1托管第一层（最接近输入）的模型参数，而worker4托管第4层（最接近输出）。“F”、“B”和“U”分别代表前向、后向和更新操作。下标表示要在哪个worker上运行操作。由于顺序依赖性，数据一次由一个worker处理，导致大量空闲时间“泡沫”。我们可以重用数据并行的思想，通过让每个工作人员一次只处理数据元素的一个子集来减少创建时间泡沫的成本，从而使我们能够巧妙地将新计算与等待时间重叠。其核心思想是通过将一个批次拆分为多个微批次，每个微批次的处理速度应该按比例更快，并且每个工作人员在下一个微批次可用时立即开始工作，从而加快流水线的实现。有了足够的微批次，workers可以在大部分时间被利用，并且在步骤的开始和结束时具有最小的“气泡”。梯度在微批次之间平均，并且仅在所有微批次完成后更新参数。模型拆分的工人数量通常称为“管道深度”。在前向传递过程中，一个工作人员只需将其层块（称为“激活”）的输出发送给下一个工作人员；在向后传递期间，它只将这些激活的梯度发送给前一个工作人员。对于如何排列这些通道以及如何跨微批次聚合梯度，存在很大的设计空间。例如，GPipe方法让每个工作进程连续向前和向后传递，然后在最后同步聚合来自多个微批次的梯度；而PipeDream安排每个工作人员交替处理的前向和后向传递。图例：GPipe和PipeDream流水线方案对比，每批使用4个微批次。微批次1-8对应于两个连续的批次数据。图中的“number”表示操作的是哪个micro-batch，下标表示workerID。请注意，PipeDream通过使用陈旧参数执行一些计算来提高效率。3TensorParallelPipelineParallelism将模型逐层“垂直”拆分，也可以在一层内“水平”拆分某些操作，通常称为张量训练。对于许多现代模型（例如Transformer），计算瓶颈是激活批处理矩阵与大权重矩阵的乘法。矩阵乘法可以被认为是行和列对之间的点积；您可以在不同的GPU上计算独立的点积，或者在不同的GPU上计算每个点积的部分并对结果求和。无论采用何种策略，我们都可以将权重矩阵拆分为大小均匀的“分片”，将每个分片托管在不同的GPU上，并使用该分片计算整个矩阵乘积的相关部分，然后再进行通信以合并结果。一个例子是Megatron-LM，它在Transformer的自注意力和MLP层中并行化矩阵乘法。PTD-P使用张量、数据和流水线并行，其流水线调度为每个设备分配多个离散层，以增加网络通信为代价减少气泡损失。有时网络输入可以跨维度并行化，与交叉通信相关的高度并行计算。序列并行是一种思想，其中输入序列在时间上被划分为子实例，通过允许计算继续进行到更细粒度的示例来按比例减少峰值内存消耗。4专家混合(MoE)使用专家混合(MoE)方法，只有网络的一小部分用于计算任何一个输入的输出。一个示例方法是拥有多组权重，网络可以在推理时通过门控机制选择使用哪一组权重，这样可以在不增加计算成本的情况下启用更多参数。每组权重都称为“专家”，网络应该学会为每个专家分配专门的计算和技能。不同的专家可以托管不同的GPU，从而提供一种明确的方法来扩展用于模型的GPU数量。图例：门控网络仅选择n个专家中的2个。GShard将MoETransformer的参数扩展到6000亿个参数，其中只有MoE层被拆分到多个TPU设备上，其他层被完全复制。SwitchTransformer通过将一个输入路由到单个专家，将模型大小扩展到具有更高稀疏性的数万亿个参数。5其他内存高效设计还有许多其他计算策略可以使训练越来越大的神经网络更容易处理。例如：要计算梯度，需要保存原始激活，这会消耗大量设备RAM。检查点（也称为激活重新计算）存储激活的任何子集，并在向后传递期间及时重新计算中间激活，以最多一次额外的完整前向传递的计算成本节省大量内存。人们还可以通过有选择地激活重新计算来不断权衡计算和内存成本，重新计算是对存储成本相对较高但计算成本较低的激活子集的检查。混合精度训练使用较低精度的数字（最常见的是FP16）来训练模型。现代加速器可以使用较低精度的数字来实现更高的FLOP计数并节省设备RAM。如果小心谨慎，生成的模型几乎不会失去任何准确性。卸载就是把不用的数据暂时卸载到CPU或者不同设备之间，需要的时候再读回来。简单的执行会大大减慢训练速度，但复杂的实现会预取数据，因此设备永远不必等待。这个想法的一个实现是ZeRO，它在所有可用硬件上划分参数、梯度和优化器状态，??并根据需要实例化它们。MemoryEfficientOptimizers提出了内存高效优化器来减少优化器维护的运行状态的内存占用，例如Adafactor。压缩也可用于在网络中存储中间结果。例如，Gist压缩了为反向传递保存的激活；DALL-E在同步之前压缩梯度。
