OpenAI研究 训练大型神经网络的技术
大型神经网络是 AI 最近取得的许多进展的核心，但训练它们是一项艰巨的工程和研究挑战，需要编排 GPU 集群来执行单个同步计算。
大型神经网络是 AI 最近取得的许多进展的核心，但训练它们是一项艰巨的工程和研究挑战，需要编排 GPU 集群来执行单个同步计算。随着集群和模型规模的增长，机器学习从业者已经开发出越来越多的技术来在许多 GPU 上并行进行模型训练。乍一看，理解这些并行技术似乎令人望而生畏，但只要对计算结构进行一些假设，这些技术就会变得更加清晰——此时，您就像网络一样在 A 到 B 之间穿梭不透明位交换数据包周围的穿梭机。
没有并行性
训练神经网络是一个迭代过程。在每次迭代中，我们都会向前传递模型的 层 ，以计算一批数据中每个训练示例的输出。然后另一遍 通过层向后进行，通过计算每个参数的梯度 来传播每个参数对 最终输出的影响程度 。批处理的平均梯度、参数和一些每个参数的优化状态被传递给优化算法，例如 Adam，它计算下一次迭代的参数（它应该对您的数据有更好的性能）和新的每参数优化状态。随着训练迭代批量数据，模型不断发展以产生越来越准确的输出。
各种并行技术将这个训练过程切入不同的维度，包括：
- 数据并行——在不同的 GPU 上运行批处理的不同子集；
- 流水线并行——在不同的 GPU 上运行模型的不同层；
- 张量并行——分解单个操作的数学运算，例如矩阵乘法，以便在 GPU 之间拆分；
- Mixture-of-Experts——只用每一层的一小部分处理每个例子。
（在这篇文章中，我们假设您正在使用 GPU 来训练您的神经网络，但同样的想法也适用于那些使用任何其他 神经网络加速器的人。）
数据并行
数据并行 训练意味着将相同的参数复制到多个 GPU（通常称为“worker”），并为每个 GPU 分配不同的示例以同时处理。单独的数据并行性仍然需要您的模型适合单个 GPU 的内存，但可以让您以存储许多重复的参数副本为代价来利用多个 GPU 的计算。也就是说，有一些策略可以增加 GPU 可用的有效 RAM，例如在两次使用之间临时将参数卸载到 CPU 内存。
当每个数据并行工作者更新其参数副本时，它们需要协调以确保每个工作者继续具有相似的参数。最简单的方法是在 worker 之间引入阻塞通信：（1）独立计算每个 worker 上的梯度；(2) 平均工人之间的梯度；(3) 对每个工人独立计算相同的新参数。步骤 (2) 是一个阻塞平均值，它需要传输大量数据（与工作人员数量乘以参数大小成正比），这可能会损害您的训练吞吐量。有多种 异步同步方案 可以消除这种开销，但它们会影响学习效率；在实践中，人们通常坚持使用同步方法。
管道并行性
通过 Pipeline Parallel 训练，我们跨 GPU 划分模型的顺序块。每个 GPU 仅包含一小部分参数，因此同一模型在每个 GPU 上消耗的内存按比例减少。
将大型模型拆分为连续层的块非常简单。然而，层的输入和输出之间存在顺序依赖性，因此在工作人员等待前一台机器的输出用作其输入时，简单的实现可能会导致大量空闲时间。这些等待时间块被称为“气泡”，浪费了空闲机器可以完成的计算。
向前
落后
更新
闲置的
我们可以重用数据并行的想法，通过让每个工作人员一次只处理数据元素的一个子集来降低泡沫成本，从而使我们能够巧妙地将新计算与等待时间重叠。核心思想是将一个批次拆分为多个微批次；每个 microbatch 的处理速度应该成比例地更快，并且每个 worker 在下一个 microbatch 可用时立即开始处理，从而加快管道执行。有了足够的微批次，workers 可以在大部分时间被利用，并且在步骤的开始和结束时气泡最小。梯度在微批次之间进行平均，并且只有在所有微批次完成后才会更新参数。
模型被拆分的工人数量通常称为 管道深度。
在前向传播过程中，工作人员只需要将其层块的输出（称为激活）发送给下一个工作人员；在向后传递期间，它只将这些激活的梯度发送给前一个工作人员。在如何安排这些遍以及如何跨微批次聚合梯度方面存在很大的设计空间。 GPipe 让每个工作进程连续向前和向后传递，然后在最后同步聚合来自多个微批次的梯度。 PipeDream 改为安排每个工作人员交替处理前向和后向传递。
向前落后更新闲置的G管道
白日梦
张量并行性
管道并行性按层“垂直”拆分模型。也可以在一层内“水平”拆分某些操作，通常称为 Tensor Parallel 训练。对于许多现代模型（例如 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 ，它将参数、梯度和优化器状态拆分到所有可用的硬件上，并根据需要实现它们。
- Memory Efficient Optimizers 已经被提出来减少由优化器维护的运行状态的内存占用，例如 Adafactor。
- 压缩 也可用于在网络中存储中间结果。例如， Gist 压缩了为反向传递保存的激活； DALL·E 在同步之前压缩梯度。
相关文章
近期评论
-
来自： GPT-4V（ision）系统卡
-
来自： GPT-4V（ision）系统卡
-
来自： OpenAI研究 用对抗样本攻击机器学习