大模型分布式训练并行技术(五)混合并行
你可能在很多地方听过3D并行技术,我们之前讨论的数据并行,流水线并行,张量并行都是属于1D并行技术。
在某些分类中,流水线并行和张量并行都被划归为模型并行技术。
混合并行技术是指同时使用多种并行技术,比如数据并行和模型并行,或者数据并行和流水线并行,或者数据并行和张量并行。
数据并行和流水线并行的结合,是一种非常常见的2D混合并行技术。
下图是来自 DeepSpeed 流水线并行教程,演示了如何将 DP 与 PP 结合起来。
这里重要的是要看到 DP Rank 0 看不到 GPU2,而 DP Rank 1 看不到 GPU3。对于 DP 来说,只有 GPU 0 和 1,它向其中提供数据,就好像只有 2 个 GPU 一样。GPU0 使用 PP“秘密”将部分负载offload到 GPU2。而 GPU1 也通过 GPU3 来做同样的事情。
由于每个维度至少需要 2 个 GPU,因此这里至少需要 4 个 GPU。
为了实现更高效的训练,就出现了3D 并行,将 PP 与 TP 和 DP 相结合。如下图所示。
该图来自博客3D parallelism: Scaling to trillion-parameter models。
由于每个维度至少需要 2 个 GPU,因此这里至少需要 8 个 GPU。
也就是常说的4D并行,ZeRO 是一个由 DeepSpeed 提出的技术,可以将模型的参数划分为多个部分,每个部分在不同的设备上进行计算,最后将结果进行汇总。本质上是ZeRO-DP 与 PP(和可选的 TP)结合。
ZeRO-DP 见[大模型分布式训练并行技术(二)Z]
注意:当 ZeRO-DP 与 PP(和可选的 TP)结合时,它通常只启用 ZeRO 第 1 阶段(优化器分片)。
虽然理论上可以将 ZeRO 第 2 阶段(梯度分片)与PP结合使用,但这会对性能产生不利影响。
每个微批次都需要一个额外的 Reduce-Scatter 集合来在分片之前聚合梯度,这可能会增加大量的通信开销。根据PP的性质,会使用较小的微批次Macro-Batch,而重点是尝试平衡算术强度(微批次大小)和最小化流水线气泡(微批次数量)。因此,这些通信成本将受到影响。
此外,由于 PP,层数已经比正常情况少,因此内存节省不会很大。PP 已经将梯度大小减少至 1/PP,因此在此基础上的梯度分片节省不如纯 DP 显著。
出于同样的原因,ZeRO 第 3 阶段也不是一个好的选择——需要更多的节点间通信。
由于我们有 ZeRO,另一个好处是 ZeRO-Offload。由于这是第 1 阶段,因此优化器状态可以卸载到 CPU。
[2] 流水线并行教程
[3] 3D parallelism: Scaling to trillion-parameter models
仓库上有原始的Markdown文件,完全开源,欢迎大家Star和Fork!