OpenAI 大型神经网络训练技术

训练大型神经网络需要协调 GPU 集群以执行单次同步计算。由于模型和集群规模的增长，实践者使用各种并行技术将计算负载和内存需求分布到多个硬件加速器上。

数据并行

数据并行通过将相同的模型参数复制到每个工作节点，并为每个节点分配不同的训练批次子集，以实现多 GPU 的使用。虽然这可以利用许多 GPU 的计算能力，但模型仍必须能够适配单个 GPU 的显存。

为了在工作节点之间保持一致性，系统必须协调参数更新。标准的同步方法包括三个步骤：

  1. 在每个工作节点上独立计算梯度。
  2. 在所有工作节点之间对梯度进行平均（阻塞通信步骤）。
  3. 在每个工作节点上独立计算新参数。

虽然存在异步同步方案以降低阻塞平均的开销，但它们往往会降低学习效率，导致大多数实践者更倾向于同步方法。

流水线并行

流水线并行将模型的顺序块（连续层）划分到不同的 GPU 上，从而降低每个设备的内存占用。

解决 “气泡” 问题

朴素的流水线实现会产生 “气泡”——即工作节点等待前一机器输出的空闲时间。为最小化这些气泡，批次被拆分为更小的微批次。每个工作节点在微批次可用时立即开始处理下一个微批次，实现计算与等待时间的重叠。梯度在微批次之间进行平均，参数更新仅在所有微批次完成后进行。

调度策略

管理前向和后向传播的两种主要调度方案：

  • GPipe：工作节点连续处理前向和后向传播，在结束时同步聚合来自多个微批次的梯度。
  • PipeDream：工作节点交替处理前向和后向传播，这可以提升效率，但可能导致部分计算使用过时的参数。

张量并行

张量并行将层内的单个操作“水平”切分到多个 GPU 上。在 Transformer 等现代架构中，主要瓶颈是将激活批矩阵与大型权重矩阵相乘。张量并行将权重矩阵划分为等大小的碎片，分别放置在不同的 GPU 上计算整体矩阵乘积的部分，然后通过通信合并结果。

具体实现包括：

  • Megatron-LM：在 Transformer 的 MLP 和自注意力层内部并行化矩阵乘法。
  • PTD-P：结合张量、数据和流水线并行，将非连续层分配给每个设备，以降低气泡开销。
  • Sequence Parallelism：将输入序列在时间维度上拆分为多个子示例，以降低峰值内存消耗。

专家混合（MoE）

专家混合（MoE）使模型能够在参数数量上扩展，而计算成本并不成比例增加。这通过门控机制实现，仅选择网络中一部分权重（即 “专家”）来计算给定输入的输出。由于不同的专家可以部署在不同的 GPU 上，MoE 为增加模型使用的 GPU 数量提供了可扩展的方式。

内存节省设计

除了并行之外，还采用多种策略来降低训练所需的设备显存：

  • Checkpointing (Activation Recomputation)：不将所有原始激活保存用于梯度计算，而是仅存储一部分，并在反向传播时即时重新计算中间激活。选择性激活重新计算进一步优化，仅对存储成本高但计算成本低的激活进行检查点。
  • Mixed Precision Training：使用低精度数值（通常为 FP16）进行训练，可提升 FLOP 计数并在几乎不损失精度的情况下节省设备显存。
  • Offloading：将未使用的数据临时移动到 CPU 或其他设备。ZeRO 实现将参数、梯度和优化器状态跨硬件拆分，并在需要时进行实例化。
  • Memory Efficient Optimizers：使用如 Adafactor 等优化器，以降低运行状态的内存占用。
  • Compression：压缩中间结果，例如 Gist 对反向传播的激活进行压缩，或 DALL·E 在同步前压缩梯度。

Sources