大规模神经网络的训练技术

# 大型神经网络的训练技术 来源：https://openai.com/index/techniques-for-training-large-neural-networks/ OpenAI 大型神经网络是最近人工智能许多进步的核心，但训练它们是一项困难的工程和研究挑战，需要协调一个GPU集群来执行单一同步计算。 大型神经网络是最近人工智能许多进步的核心，但训练它们是一项困难的工程和研究挑战，需要协调一个GPU集群来执行单一同步计算。随着集群和模型规模的增加，机器学习从业者已经开发了越来越多的技术来在多个GPU上并行化模型训练。乍一看，理解这些并行化技术可能显得令人望而却步，但只要对计算结构做出一些假设，这些技术就变得清晰多了——此时，你所做的就是在A和B之间传输不透明的数据，就像网络交换机传输数据包一样。 训练神经网络是一个迭代过程。在每次迭代中，我们通过模型的各层(https://developer.nvidia.com/blog/deep-learning-nutshell-core-concepts/#layer)进行前向传播，为数据批次中的每个训练示例计算输出。然后另一次传播向后(https://youtu.be/Ilg3gGewQ5U)通过各层进行，通过计算关于每个参数的梯度(https://youtu.be/IHZwWFHWa-w)来传播每个参数对最终输出的影响程度。批次的平均梯度、参数和某些每参数优化状态被传递给优化算法（如Adam(https://arxiv.org/abs/1412.6980)），该算法计算下一次迭代的参数（应在你的数据上有略微更好的性能）和新的每参数优化状态。当训练遍历数据批次时，模型不断演变以产生越来越准确的输出。 各种并行化技术沿着不同的维度对这个训练过程进行切分，包括： - 数据并行——在不同的GPU上运行批次的不同子集； - 管道并行——在不同的GPU上运行模型的不同层； - 张量并行——将单个操作（如矩阵乘法）的数学运算分割到多个GPU上； - 混合专家——每个示例仅由每层的一部分处理。 （在这篇文章中，我们假设你使用GPU来训练神经网络，但相同的思想适用于使用任何其他神经网络加速器(https://www.synopsys.com/ai/what-is-an-ai-accelerator.html)的人。） **数据并行**训练意味着将相同的参数复制到多个GPU（通常称为"工作进程"），并同时为每个GPU分配不同的示例进行处理。数据并行本身仍然需要你的模型能够装入单个GPU的内存，但让你能够利用多个GPU的计算能力，代价是存储参数的多个重复副本。也就是说，有一些策略可以增加GPU可用的有效RAM，例如在使用期间暂时将参数卸载到CPU内存。 当每个数据并行工作进程更新其参数副本时，它们需要进行协调以确保每个工作进程继续拥有相似的参数。最简单的方法是在工作进程之间引入阻塞通信：（1）在每个工作进程上独立计算梯度；（2）跨工作进程平均梯度(https://tech.preferred.jp/en/blog/technologies-behind-distributed-deep-learning-allreduce/)；（3）在每个工作进程上独立计算相同的新参数。第（2）步是一个阻塞平均，需要传输相当大量的数据（与工作进程数量乘以参数大小成正比），这可能会损害你的训练吞吐量。有各种异步同步方案(https://arxiv.org/abs/1106.5730)来消除这种开销，但它们会损害学习效率；实际上，人们通常坚持同步方法。 通过**管道并行**训练，我们将模型的顺序块分区到各个GPU上。每个GPU仅保持一部分参数，因此相同的模型在每个GPU上消耗的内存成比例地减少。 将大型模型分割成连续层的块是很直接的。然而，层的输入和输出之间存在顺序依赖关系，所以天真的实现可能导致大量的空闲时间，当一个工作进程等待来自前一个机器的输出作为其输入时。这些等待时间块被称为"气泡"，浪费了空闲机器本可以完成的计算。 我们可以重用数据并行中的思想来减少气泡的代价，通过让每个工作进程每次只处理数据元素的子集，允许我们巧妙地将新计算与等待时间重叠。核心思想是将一个批次分割成多个微批次；每个微批次的处理速度应该相对更快，每个工作进程在下一个微批次可用时立即开始处理它，从而加快管道执行。有足够的微批次后，工作进程可以大部分时间处于运行状态，在步骤的开始和结束处仅有最小的气泡。梯度在微批次之间进行平均，参数的更新仅在所有微批次完成后发生。 模型分割的工作进程数通常称为**管道深度**。 在前向传播期间，工作进程只需要将其层块的输出（称为激活）发送给下一个工作进程；在反向传播期间，它只发送这些激活上的梯度给前一个工作进程。有许多设计空间可以安排这些传播以及如何跨微批次聚合梯度。GPipe(https://arxiv.org/abs/1811.06965)让每个工作进程连续处理前向和反向传播，然后在最后同步聚合来自多个微批次的梯度。PipeDream(https://cs.stanford.edu/~matei/papers/2019/sosp_pipedream.pdf)则是安排每个工作进程交替处理前向和反向传播。 管道并行按层"垂直"分割模型。也可以"水平"分割层内的某些操作，这通常称为**张量并行**训练。对于许多现代模型（如Transformer(https://jalammar.github.io/illustrated-transformer/)），计算瓶颈是将激活批处理矩阵与大权重矩阵相乘。矩阵乘法(https://en.wikipedia.org/wiki/Matrix_multiplication)可以看作是行和列对之间的点积；可以在不同的GPU上计算独立的点积，或在不同的GPU上计算每个点积的部分并求和结果。使用任一策略，我们可以将权重矩阵切分成均匀大小的"分片"，在不同的GPU上托管每个分片，并使用该分片来计算总体矩阵乘积的相关部分，然后进行通信来合并结果。 一个例子是Megatron-LM(https://nv-adlr.github.io/MegatronLM)，它在Transformer的自注意力和MLP层内并行化矩阵乘法。PTD-P(https://arxiv.org/abs/2104.04473)使用张量、数据和管道并行；其管道调度为每个设备分配多个非连续层，以降低气泡开销为代价的更多网络通信。 有时网络的输入可以沿着具有高度并行计算相对于交叉通信的维度进行并行化。序列并行(https://arxiv.org/abs/2205.05198)是一个这样的想法，其中输入序列在时间上分割成多个子示例，通过允许计算以更细粒度大小的示例进行，按比例减少峰值内存消耗。 对于**混合专家（MoE）**(https://arxiv.org/abs/1701.06538)方法，网络的只有一部分被用来计算任何一个输入的输出。一个示例方法是拥有许多权重集合，网络可以通过在推理时的门控机制选择使用哪个集合。这使得可以拥有更多参数而不增加计算成本。每个权重集合称为"专家"，希望网络会学会为每个专家分配专门的计算和技能。不同的专家可以托管在不同的GPU上，提供了一个清晰的方式来扩展模型使用的GPU数量。 混合专家（MoE）层的图示。仅由门控网络选择n个专家中的2个。（图片改编自：Shazeer等人，2017） 有许多其他计算策略可以使训练越来越大的神经网络更具可操作性。例如： - 要计算梯度，你需要保存原始激活，这可能消耗大量设备RAM。**检查点**(https://arxiv.org/abs/1604.06174)（也称为激活重计算）存储任何激活子集，并在反向传播期间及时重计算中间激活。这以最多一次额外完整前向传播的计算成本节省了大量内存。也可以通过选择性激活重计算(https://arxiv.org/abs/2205.05198)在计算和内存成本之间不断权衡，这是对相对存储成本更高但计算成本更低的激活子集进行检查点。 - **混合精度训练**(https://arxiv.org/abs/1710.03740)是使用较低精度数字（最常见的是FP16(https://en.wikipedia.org/wiki/Half-precision_floating-point_format)）来训练模型。现代加速器可以用较低精度数字实现更高的FLOP计数，你也可以节省设备RAM。只要适当小心，生成的模型几乎不会失去准确性。 - **卸载**是暂时将未使用的数据卸载到CPU或不同设备之间，稍后在需要时读回。天真的实现会大大减慢训练速度，但复杂的实现会预取数据以便设备永远不需要等待。此想法的一个实现是ZeRO(https://arxiv.org/abs/1910.02054)，它在所有可用硬件上分割参数、梯度和优化器状态，并根据需要实现它们。 - **内存高效的优化器**已被提出来减少优化器维持的运行状态的内存占用，例如Adafactor(https://arxiv.org/abs/1804.04235)。 - **压缩**也可用于存储网络中的中间结果。例如，Gist(https://www.microsoft.com/en-us/research/uploads/prod/2018/04/fiddle-gist-isca18.pdf)压缩为反向传播而保存的激活；DALL·E(https://openai.com/index/dall-e/)在同步梯度之前压缩梯度。 在OpenAI，我们正在从底层基础设施一直到为现实世界问题部署它们来训练和改进大型模型。如果你想把这篇文章中的思想付诸实践——特别是与我们的Scaling和Applied Research团队相关的——我们正在招聘(https://openai.com/careers/research-engineer/)！