LLM学习之路-01-第一章-预训练/大模型分布式训练并行技术（五）混合并行

大模型分布式训练并行技术（五）混合并行

引言

你可能在很多地方听过3D并行技术，我们之前讨论的数据并行，流水线并行，张量并行都是属于1D并行技术。

在某些分类中，流水线并行和张量并行都被划归为模型并行技术。

混合并行技术是指同时使用多种并行技术，比如数据并行和模型并行，或者数据并行和流水线并行，或者数据并行和张量并行。

DP+PP

数据并行和流水线并行的结合，是一种非常常见的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。

DP+PP+TP

为了实现更高效的训练，就出现了3D 并行，将 PP 与 TP 和 DP 相结合。如下图所示。

该图来自博客3D parallelism: Scaling to trillion-parameter models。

由于每个维度至少需要 2 个 GPU，因此这里至少需要 8 个 GPU。

DP+PP+TP+ZeRO

也就是常说的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。