Distributed Training/zh
| Article | |
|---|---|
| Topic area | Deep Learning |
| Prerequisites | Stochastic Gradient Descent, Backpropagation |
概述
分布式训练是将单个神经网络训练作业拆分到多个处理器、加速器或机器上,使工作负载并行运行的实践。它解决了单设备训练的两个限制:模型可能太大,无法装入单个设备的内存;以及在大数据集上收敛所需的实际时间可能令人无法接受。具有数千亿参数的现代基础模型如果没有它根本无法训练,即使是较小的模型也通常使用它来将训练时间从数周缩短到数小时。该领域结合了数值优化、高性能计算和系统工程的思想:如何在工作节点之间协调随机优化,如何在设备之间高效地移动张量,以及如何调度计算以使加速器保持忙碌。
两个基本的轴是数据并行,其中每个工作节点持有模型的副本并处理批次的不同分片;以及模型并行,其中模型本身被拆分到工作节点之间。现实世界的系统将两者与流水线并行和优化器状态分片相结合,从而扩展到数千个加速器,同时保持较高的算术强度和可承受的通信开销。
动机
三种压力将训练推离单一设备。首先是参数数量:具有 $ N $ 个参数的 16 位 精度模型至少需要 $ 2N $ 字节用于权重,梯度和优化器状态也需要类似的量,因此一个 700 亿参数的模型在混合精度下,一旦加入 Adam 矩,就轻易超过 1 TB 内存。其次是数据量:在单设备吞吐量下用数万亿个 token 训练一个大型语言模型需要数年时间。第三是样本效率:更大的有效批次大小通过数据并行变得 made 可行,可以稳定梯度并改善泛化,直到临界批次大小,超过该值后回报会递减。
然而,扩展并非免费。工作节点之间的通信与计算相互竞争,同步屏障可能使设备闲置。分布式训练的核心工程问题是如何在保留原始算法优化动态的同时,重叠、压缩或消除该通信。
数据并行
在同步数据并行训练中,$ K $ 个工作节点中的每一个都持有参数 $ \theta $ 的相同副本,并处理大小为 $ B $ 的不同微批次。每个工作节点计算局部梯度 $ g_k = \nabla_\theta L_k(\theta) $,然后工作节点通过 all-reduce 操作交换梯度,使每个工作节点都持有平均值 $ \bar{g} = \frac{1}{K}\sum_{k=1}^{K} g_k $,并且每个工作节点在本地应用相同的优化器更新:
$ {\displaystyle \theta \leftarrow \theta - \eta \, \bar{g}.} $
由于所有副本以相同状态开始,并对相同的平均梯度应用相同的更新,它们保持逐位一致(除归约中的非确定性外)。有效批次大小为 $ K B $,这通常需要学习率预热和线性缩放规则来保持稳定性。
all-reduce 原语可用环形或树形算法实现;ring all-reduce每个工作节点传输梯度大小的 $ 2(K-1)/K $ 倍,在均匀互联上是带宽最优的。NCCL、Gloo 和 MPI 都提供了经过调优的实现。
一种较旧的替代方案是参数服务器,即由一个中心进程持有规范的参数并聚合工作节点的更新。异步变体允许工作节点无需等待即可推送陈旧的梯度,以优化质量换取吞吐量;这种方法在密集深度学习中大多已成为历史,但在某些推荐器和嵌入工作负载中仍然存在,因为稀疏性和负载不平衡有利于这种方法。
模型并行
当单层的参数或激活超过设备内存时,该层本身必须被分区。在张量并行中,矩阵乘法 $ Y = X W $ 沿隐藏维度被拆分:每个工作节点持有列条 $ W_k $,计算 $ Y_k = X W_k $,部分输出根据层类型进行拼接或求和。Megatron-LM 推广了这种方案用于transformer 块,通过对注意力和 MLP 投影进行分片,使得每个块在前向和反向传递中只需要两次all-reduce。
张量并行在一个节点内扩展良好,在该节点内 NVLink 等高带宽链路连接着一小组 GPU(通常为 4 到 8 个)。在跨节点的情况下,每步通信成本快速增长,通常与其他形式的并行结合使用,而不是单独使用。
流水线并行
流水线并行按深度将网络在工作节点之间分区:工作节点 $ k $ 持有连续的层块,并将激活转发到工作节点 $ k+1 $。简单地一次运行一个微批次会使大多数工作节点空闲(称为"流水线气泡")。GPipe 通过将小批次拆分为 $ M $ 个微批次并重叠它们的前向传递来解决这个问题;对于 $ P $ 个阶段,气泡比例大致为 $ (P-1)/(M+P-1) $。PipeDream-Flush 和 1F1B 调度更激进地交错前向和反向传递,以减少激活内存。
流水线并行对负载均衡敏感:单个慢速阶段会停滞整条流水线。从业者会分析每个阶段的耗时并调整层的分配,有时会在各阶段之间转移工作或复制开销小的层。
混合并行与 3D 并行
在集群规模下,任何单一轴都不够用。常见的安排——有时称为 3D 并行——组合了数据并行、张量并行和流水线并行。张量并行处理节点内最大的层,流水线并行将模型在节点之间分区,而数据并行在副本组之间复制整条流水线以获得吞吐量。序列并行沿序列维度对激活进行分片,是用于超长上下文的第四个轴。
选择分区是一个离散优化问题,由内存预算、互联拓扑以及模型的注意力和前馈形状决定。Alpa、GSPMD 以及 DeepSpeed 和 Megatron 中的分区例程会自动搜索这一空间。
ZeRO 与分片优化器状态
Zero Redundancy Optimizer(ZeRO)通过对每个副本存储的内容进行分片来降低数据并行的内存开销。ZeRO-1 将优化器状态在数据并行组之间分区;ZeRO-2 也对梯度进行分区;ZeRO-3(也称为 Fully Sharded Data Parallel,FSDP)则对参数本身进行分区,在每一层按需聚合。由于 gather 和 reduce-scatter 原语可以与反向计算组合,通信可以在很大程度上与算术运算重叠,从而以数据并行更简单的编程模型实现接近模型并行的内存占用。
结合混合精度存储和 CPU 或 NVMe 卸载,ZeRO 已使得在商品级规模的 GPU 集群上训练万亿参数的模型成为可能。
通信与同步
分布式训练中占主导地位的开销是移动张量。有几种技术可以减少或隐藏这一开销:
- 梯度压缩,包括 1-bit SGD、signSGD 和 PowerSGD,以引入必须加以控制的偏差或方差为代价,减少了每次 all-reduce 的字节数。
- Local SGD 在每个工作节点上执行 $ H $ 步后再求平均,以略微的优化漂移为代价降低了通信频率。
- 计算与通信的重叠利用了这样一个事实:早期层的梯度在后续层仍在计算时就已经就绪;PyTorch DDP 和 Horovod 等框架以桶为单位调度异步 all-reduce,使得网络在反向传递期间始终在传输数据。
- 拓扑感知的集合通信优先使用节点内高带宽链路进行早期归约,只在最后才跨越慢速链路。
实践考量与局限
分布式训练会放大每一个小故障。在数千个加速器规模下,落后节点、链路错误和静默数据损坏变得司空见惯,因此检查点频率、确定性重启和弹性重分片都需要大量工程关注。诸如损失尖峰不稳定之类的数值问题在大批次大小下更为常见,可能需要梯度裁剪、更低的学习率,或从最近的检查点恢复。
并行策略的选择很少具有可移植性:为某种集群拓扑调优的分区在另一种拓扑上可能远非最优,改变模型形状可能需要重新分片优化器状态。通信库、内核启动开销和每设备的内存碎片都成为头等关注的问题。尽管存在这些工程负担,分布式训练仍是通向当今前沿规模模型的唯一已知途径,集合通信算法、分片策略和加速器互联方面的改进将继续拓宽其实际可行的边界。