论文地址-《ZeRO++: Extremely Efficient Collective Communication for Giant Model Training》

摘要

• ZeRO的缺点： 当在低带宽集群上进行训练时，或者在规模上迫使每个 GPU 的批量大小变小时，ZeRO 的有效吞吐量受到限制，因为在前向传递、后向传递和平均梯度中收集权重的通信量很大。
• 本文贡献：本文介绍了三种通信量缩减技术， 我们统称为 ZeRO++，针对 ZeRO 中的每个通信集合。
  • 首先是基于块量化的全收集。
  • 其次是数据重新映射，以牺牲通信换取更多内存。
  • 第三是一种新颖的基于全对全的量化梯度平均范式，作为减少分散集合的替代，尽管传输低精度数据，但仍保持准确性。

Ⅰ 介绍

• 训练大规模模型需要有效地使用数百甚至数千个 GPU 设备上的聚合计算能力和内存。有两种流行的方法，即 3D 并行和ZeRO。
  • 3D并行即数据并行+流水线并行+张量并行，主要缺点是系统和用户的复杂性。它给数据科学家和人工智能从业者带来了重构单个 GPU 代码以实现 3D 并行性的负担， 这非常重要而且通常很麻烦。
  • ZeRO 提供了一种无需重构模型代码的替代方案。 ZeRO 是数据并行性的内存高效变体 [2, 6]，其中模型状态在所有 GPU 上进行分区，而不是进行复制，并在训练期间使用基于收集的即时通信集合进行重建。

1.1 ZeRO的局限性

• ZeRO的局限性
  • 在两个关键场景中，ZeRO 的效率可能会因通信开销而受到限制：
    • 低带宽集群：如图 1(a) 所示，低带宽集群上的每个 GPU 吞吐量仅为高带宽集群的一半。
    • 每个 GPU 的批量大小非常小：由于全局批量大小无法无限增加而不减慢模型收敛速度，因此在数千个 GPU 上进行训练迫使每个 GPU 的批量大小非常小，从而降低了计算与通信的比率，从而产生通信瓶颈。如图 1(b) 所示，每个 GPU 的吞吐量受到每个 GPU 的小批量大小的严重影响。
  • 前人工作的局限：很少研究关注优化 ZeRO 的端到端通信效率方面。

1.2 ZeRO++

• 分析ZeRO的通信量

  • 在前向传递过程中，ZeRO [29] 进行 all-gather 操作来收集所有模型层训练所需的所有参数（M）
  • 在后向传递中，ZeRO 首先使用 all-gather 重新收集参数（M），然后每个 GPU 可以计算局部梯度。之后，ZeRO 运行reduce-catter 函数以在加速器之间聚合和重新分配梯度（M）
  • 总的来说，ZeRO 的总通信量为3 M，均匀分布在 2 个 all-gather 和 1 个 reduce-catter 上。

• 为了减少这些通信开销，ZeRO++针对上述三个通信集合做了三组通信优化

  • ZeRO 的量化权重通信（qwZ）：为了减少前向传播期间 all-gather 的通信量，可以使用量化，但是为了保持良好的模型训练精度，我们采用基于块的量化，对模型参数的每个子集进行独立量化。目前还没有高性能基于块的量化的实现。因此，我们从头开始实现高度优化的量化 CUDA 内核。
  • ZeRO 的分层权重分区 (hpZ) ：为了减少后向过程中权重 all-gather 的通信开销，我们用 GPU 内存来交换通信。我们不是在所有机器上分散整个模型权重，而是在每台机器内维护完整的模型副本。 以更高的内存开销为代价，这使我们能够用机器内 all-gather 代替昂贵的跨机器权重全收集，由于更高的机器内通信带宽，速度要快得多。
  • ZeRO 的量化梯度通信 (qgZ)
    • 使用reduce-scatter来降低梯度的通信成本更具挑战性，直接应用量化来减少通信量是不可行的。主要问题是，即使通过结合基于块的量化来减少分散操作，它仍然会严重损害模型训练的准确性。背后的关键原因是量化会降低值的精度，低精度值会累积并放大误差。
    • 因此，我们提出了一种新颖且更高效的梯度通信范例作为减少 reduce-scatter 的一般替代品，其中在通信过程中使用基于块的 INT4 量化来压缩梯度以减少通信量，但在之前恢复了完整的精度归约算子以保持训练精度。
    • 我们称之为 qgZ ，其设计目的是：1) 克服如果我们简单地在 INT4/INT8 中实现reduce-scatter，则因低精度归约而导致的显着准确度损失，以及 2) 避免长时间的准确度下降和显着的延迟开销。基于环或树的reduce-scatter 所需的量化和反量化步骤序列，即使我们以全精度进行缩减。
    • 此外，qgZ 利用现代 GPU 集群的分层性质，其中节点内带宽明显高于节点间，首先减少节点内的梯度，然后再进行跨节点减少，以最大限度地减少节点间通信量，从而导致 与 FP16 减少分散相比，通信量减少了 2/4 倍 (INT8/4)。我们通过管道化节点内和节点间通信并进行 CUDA 内核融合，进一步减少 qgZ 的端到端延迟。

• 通信量减少：通过整合上述三个组件，我们将跨节点通信量从3M减少到0.75M。
  • 前向 all-gather期间，从M变成0.5M
  • 后向 all-gather 期间，从M减少到0
  • 通过将梯度上的后向fp16 reduce-scatter替换成 all-to-all INT4 reduce-scatter，将跨节点通信从M减少到了0.25M

Ⅱ 相关工作

• 量化
• 梯度量化
• ZeRO通信减少
  • 为了减少昂贵的跨节点通信，最近对 ZeRO-3 的优化（例如 MiCS [40]）用设备上的内存来进行通信。在 MiCS 中，GPU 集群被分为子组，模型状态在子组内划分，但在子组之间复制。通过保持较小的子组大小，MiCS 可以利用高带宽节点内互连，或使用分层通信来降低通信量。
  • ZeRO++ 中的 hpZ 采用了类似的方法，通过交换内存来减少通信。主要区别在于 hpZ 仅对权重进行辅助分区，同时保持所有其他模型状态在所有 GPU 上分区。这使得 hpZ 能够显着减少通信量，而无需 MiCS 的大量内存开销。

Ⅲ 设计

3.1 Quantized Weight Communication for ZeRO (qwZ )

• 为了减少前向全收集权值的通信开销，qwZ 在全收集过程中将 FP16 权重量化为 INT8，并在接收端将其反量化回 FP16，然后进行层计算。

• 虽然这将全收集的通信量减少了 2 倍，但这样做会导致两个主要问题：

  • 精度的降低会导致训练期间精度显着下降，如 2.3 中所述。
  • 量化和反量化开销否定通信量减少带来的任何吞吐量增益。我们在第 4 节中讨论了 qwZ 的优化实现，以最大限度地减少量化和反量化开销。

• qwZ 使用基于分块的量化来提高量化精度。如图 2 所示，每个权重张量被分成更小的块，并使用独立的量化缩放系数通过对称量化转换为 INT8。通过保持较小的量化粒度，我们可以显着缩小数字范围和粒度之间的差距。

  • 图 2(a) 中展示了执行基于块的量化与非块量化基线的量化误差的示例。
  • 图 2(b) 显示了 BERT 模型上权重量化的案例研究，其中基于块的量化将量化误差降低了 3 倍。

3.2 Hierarchical Partitioning for ZeRO (hpZ )

• 我们注意到可以针对不同的模型状态进行不同的分区，从而将通信集合限制为 GPU 的子集。鉴于在现代 GPU 集群上，节点内通信带宽明显高于节点间通信带宽，这为减少节点间通信提供了机会。
• 更具体地说，在 hpZ 中，我们通过在每个节点内保留辅助 FP16 权重分区来消除向后传递过程中的节点间全收集。我们通过创建一个由两个分区组成的分层分区策略来实现这一点：首先，所有模型状态都在所有设备上进行全局分区，如 ZeRO-3 中，我们将其称为主分区。其次，FP16 参数的辅助副本在子全局级别（例如计算节点，见图 3）进行分区，我们将其称为辅助分区。 FP16 参数的辅助副本跨多个辅助分区进行复制。
  • 考虑一个 64 节点集群，每个节点有 8 个 GPU。模型权重分为两个阶段：i）跨所有 512 个 GPU，我们称之为主分区，ii）相同的权重也在计算节点内跨 8 个 GPU 进行分区，我们称之为辅助分区。在此示例中，对于辅助分区，集群中的每个计算节点都有一份FP16权重的完整副本，划分存储在该节点的8个GPU上，因此总共有 64 个这样的副本。

3.2.1 使用 hpZ 进行训练迭代。

• 在训练迭代的前向传递过程中，我们基于所有 GPU 上的主分区收集权重。然而，一旦在前向传递过程中消耗了权重，它们就会根据辅助分区进行分区。考虑到前向和后向传递之间模型参数的时间一致性，当在后向传递期间再次需要权重时，我们都基于该辅助分区组来收集权重。 请注意，当辅助分区设置为计算节点时，这可以避免此全收集的任何节点间通信。最后，在迭代结束时，在优化器步骤期间，所有模型状态以及 fp16 参数的主副本都会根据主分区进行更新。
• 默认情况下，hpZ 辅助分区是基于节点的，但可以根据需要进行扩展以支持多个计算节点。

3.2.2 内存使用分析

• hpZ 用内存来换取通信效率。分析这种权衡很重要。
• 图 4 提供了 100B 参数大小的典型大型语言模型的具体内存使用估计，主要组大小为 1024 个 GPU，次要组大小为 16 个 GPU（例如 DGX-2 V100 节点）。
  • 使用我们提出的方法，hpZ 消耗的内存比 ZeRO-3 多 8.9 倍，但我们的方法的内存需求仍然比标准 DP 少 114 倍。
  • 内存使用量的这种边际增加可以通过高效的节点内通信调度来补偿。通过消除或减少反向传递的节点间通信，hpZ 将 ZeRO 的端到端通信减少了 1.5 倍，同时仍然支持数千亿参数的模型训练。

3.3 Quantized Gradients Communication for ZeRO (qgZ )

• 在本节中，我们提出了一种基于 all-to-all 集体的新颖的量化 reduce-scatter 算法 qgZ，该算法通过用 INT4 量化数据替换 FP16，使梯度 reduce-scatter 的通信量减少 4 倍，同时克服了第 1 节中描述的精度损失挑战，以及我们将在本节中概述的众多系统挑战。
• qgZ 利用 all-to-all 集合体来实现量化 reduce-scatter ，其中包括三个主要组成部分：
  • 1）基于 all-to-all 的量化梯度reduce-scatter实现
  • 2）通过分层集合体减少通信量
  • 3）张量切片重新排序正确的渐变放置

3.3.1 All-to-all 基本实现

• 原始的方法使用基于环的 reduce-scatter

  • 实现量化 reduce-scatter 的一种简单方法是直接将量化和反量化应用于基于环的 reduce-scatter ，同时避免因reduce 造成的精度损失，如图5左侧所示。我们可以在每个GPU上注入量化和反量化。一旦GPU从前序收到，我们对其进行反量化以恢复完全精度并进行局部缩减。接下来我们可以量化局部缩减输出并将量化数据传递给其后继者。为了完成整个reduce-scatter，连续量化和反量化内核的数量等于正在使用的GPU的数量（即n）。
  • 缺点：在现有的基于环的减少分散集合上应用量化和反量化将由于多个连续的量化和反量化步骤而导致高通信延迟和低值精度。尽管最近的基于树的集合（如 Blink[38]）可以将顺序内核的数量从 n 减少到 log(n)，但长延迟和低精度问题尚未完全解决。
  
  

• 解决办法

  • 采用 1 跳全对 all-to-all 进行梯度通信。
  • 如图 5 右侧所示，我们首先对给定张量应用量化，然后在所有 GPU 之间进行 all-to-all 通信。总而言之，我们应用另一次反量化来恢复数据精度，然后对高精度值进行缩减以获得最终的梯度缩减输出。通过用我们的all-to-all集合替代基于环的解决方案，我们将顺序量化+反量化内核的数量从GPU的数量减少到1。因此，我们解决了在reduce中应用量化时的长延迟和低精度问题。

3.3.2 减少节点内的通信量

• 用 all-to-all 引入了一个新问题：尽管数据量化了，但是节点间通信量增加了。如下图所示，每个机器都会产生 ​N*M/Z 的跨节点通信量，远大于reduce-scatter。
• 解决办法
  • 采用分层2跳 all-to-all，而不是1跳。首先是节点内 all-to-all，其次是节点间 all-to-all，下面第二张图所示。
  • 首先利用机器内部GPU之间的高带宽链路，对量化数据进行节点内 all-to-all，然后对数据进行反量化，并对反量化后的数据进行 reduce。经过节点内量化、all-to-all、反量化和缩减之后，我们将每个 GPU 的数据大小从 ​M/Z 减少到 ​M/(Z *N )。
  • 接着，我们进行节点间all-to-all通信，类似于我们上面描述的1跳all-to-all。假设现在每个 GPU 只需要发送 ​M/(Z * N ) 数据，则每台机器的通信量现在为 ​M/(Z * N ) * N = M/Z 。
• 通过采用这种分层的全对所有通信作为 2 跳方法，我们完美地解决了 1 跳方案中的通信量爆炸问题。请注意，即使总通信量增加了一倍（一个节点内，另一个节点间），考虑到 NVLink/NVswitch 高带宽，节点内通信引入的开销可以忽略不计，并且跨节点流量已显着减少，这是主要瓶颈在梯度通信中。

3.3.3 张量切片重排以正确放置数据

• 使用2跳 all-to-all，节点间通信量符合预期，但这引入了梯度错位问题。我们使用 2x2 示例来描述此问题，其中我们有 2 台机器，每台机器有 2 个 GPU。如图 8 所示，正确的最终梯度放置在图中显示为绿色框，其中 GPU 0 保存最终梯度分区 1，GPU 1 保存梯度分区 2，依此类推。

• 解决办法
  • 通过张量切片重新排序来解决这个问题。如图9所示，在节点内all-to-all开始之前，我们首先交换切片2和3的张量切片顺序，如橙色箭头所示。然后，在节点内all-to-all完成后，G1现在具有梯度2，而G2具有梯度3。
  • 基于方程 1 和方程 2，我们可以将每个原始张量切片位置（即方程 1）映射到每个 GPU 上的新张量切片位置（即方程 2），以纠正最终梯度错位问题。

3.3.4 ZeRO++通信量分析

Ⅳ 优化实现

• 在本节中，我们讨论两个关键优化，使 ZeRO++ 能够充分实现 4 倍通信量减少的潜力，从而提高吞吐量，而不受实现开销的限制：
  • 重叠不同的通信和计算流，这样做可以实现更好的资源利用率。
  • 针对量化、反量化和张量切片重新排序运算符优化了 CUDA 内核，并在适当时跨这些运算符进行内核融合，以最大限度地减少内存流量开销。

4.1 重叠计算和通信

为了减少端到端通信时间，我们将量化计算与通信重叠，以在前向和后向传递中全部收集权重。对于基于梯度减少分散的分层实现，我们将节点内通信与节点间通信重叠。

4.1.1 通信计算在权重上重叠

• 我们使用两个关键功能实现通信计算重叠：
  • i）我们跟踪模型层的执行顺序以获取它们将被获取的顺序。
  • ii) 我们保证异步量化执行。
• 有了这两个特性，作为每一层的ZeRO获取参数，当前层的通信和下一层的量化可以在不同的CUDA流上同时启动。

4.1.2 用于梯度通信的分层集体Hierarchical Collectives

• 我们基于all-to-all的梯度通信分为两个阶段：首先是节点内通信，然后是节点间通信。节点间通信取决于节点内通信的结果，因此，在简单的实现中，节点间链路在节点内通信期间是空闲的。
• 为了通过并行利用节点间和节点内链接来减少延迟，我们对输入梯度张量和节点内通信和节点间通信之间的管道传输进行分块。

• 正确执行此管道会对我们的张量切片重新排序过程产生影响。我们拥有的管道阶段越多，重新排序所需的细粒度张量切片就越多。 因此，我们还提出了一种广义张量切片重新排序方案作为算法2，它涵盖了w/和w/o流水线数据传输情况。这里的stages指的是我们拥有的管道阶段的数量，nodeSize是每个节点的GPU数量，nodes是节点的数量。

4.2 CUDA 内核

略

Ⅴ 评估

略