【论文笔记】【存储】Optimus-CC: Efficient Large NLP Model Training with 3D Parallelism Aware Communication C~

本文链接：https://blog.csdn.net/qq_43207982/article/details/134342599

[distributed training framework] Optimus-CC: Efficient Large NLP Model Training with 3D Parallelism Aware Communication Compression. Jaeyong Song, Jinkyu Yim, Jaewon Jung, Hongsun Jang, Hyung-Jin Kim, Youngsok Kim, Jinho Lee. ASPLOS’23

论文地址: https://dl.acm.org/doi/pdf/10.1145/3575693.3575712

用于大型NLP模型的快速可扩展分布式（3D）训练框架：Optimus-CC

摘要阅读

现代自然语言处理（NLP）模型的训练通常会采用3D并行性技术，将模型分割到多个图形处理器（GPU）上。然而，这种做法存在一个问题，即节点间通信的高昂成本。虽然通信压缩可以减少这种成本，但现有的压缩技术有其限制，并可能损害模型质量。

为解决这个问题，我们提出了一种名为Optimus-CC的分布式训练框架。该框架的优势在于：

我们不仅压缩了数据并行性的通信，还压缩了管道并行（阶段间）的通信，包括阶段间的反向传播和嵌入同步。
我们提出了新的技术，可以防止由于压缩而导致的模型质量下降。
我们对管道进行了深入分析，并优化了位于关键路径上的通信压缩，进一步减少了压缩错误。

测试结果表明，Optimus-CC可以在不牺牲模型质量的前提下，显著提高分布式训练的速度。

3D 并行：data, pipeline, tensor =》数据并行模型并行流水线并行有的文章中也会提到模型并行（MP）操作符并行（OP）专家并行（EP）反正就是各种并行

前言阅读

最近的分布式训练框架建议使用管道并行和张量并行将模型分割成多个部分。这三种类型（即，数据、管道、张量）的并行性被统称为分布式训练的3D并行性。

问题

只针对数据并行通信的现有压缩方法对于大型DL模型的分布式训练效率低下，尤其是对于使用三维并行性的模型。
- 数据并行通信不再是节点间通信的唯一来源。
- 对最近的大型NLP模型直接应用现有的压缩技术会导致模型质量（即，下游任务准确性）显著下降。
- 将压缩应用于新引入的管道并行通信会导致更多的质量下降。

这表明需要新的技术来抑制压缩错误，或者在不损失的情况下减少通信量。

现有的压缩方法忽视了来自管道调度的机会。当训练过程被管道化时，大部分的通信延迟被后续微批量的计算（即，前向和后向传递）所隐藏。因此，盲目地压缩所有的通信流量只会产生更多的压缩错误，而不会产生任何吞吐量的增益。

标题中的CC的意思是压缩通信的意思 compress communication.

Optimus-CC采用了以下三个关键思想：

压缩反向传播针对管道并行的阶段间通信。
融合嵌入同步将嵌入层的两个全归约通信融合为一个全归约通信。
选择性阶段压缩针对数据并行通信，但将压缩目标限制在几个阶段

Optimus-CC的设计目标是在不牺牲模型质量的情况下，利用三维并行中节点间通信的特性来实现吞吐量的增益。我们的方法使用了压缩反向传播、融合嵌入同步和选择性阶段压缩这三种策略

通过在管道并行的阶段间通信中应用压缩反向传播，我们可以在不影响模型质量的情况下，提高加速度。这种方法专注于管道结尾，且以惰性的方式在迭代内传播压缩错误
融合嵌入同步的策略将嵌入层的两个全归约通信融合为一个全归约通信。由于网络的开始和结束都共享一个嵌入层，因此我们发现，将同步操作融合可以在不改变数学结果的情况下，减少通信量。
选择性阶段压缩针对数据并行通信，但只将压缩目标限制在几个阶段。由于数据并行通信在对应阶段的反向传播完成后立即开始，因此，早期阶段可能会将数据并行通信放在关键路径上，所以选择性阶段压缩选择不压缩某些阶段，以减少错误。

贡献总结

总的来说，贡献可以总结如下：

提出了Optimus-CC，这是一个快速且可扩展的分布式训练框架，它在保持模型质量的同时，大幅度压缩了节点间的通信。据本文所知，Optimus-CC是首个使用节点间通信压缩来加速大规模NLP模型训练的工作。
我们提出了三种专门用于减少3D并行通信量的技术：压缩反向传播、融合嵌入同步和选择性阶段压缩。它们显著地提高了训练吞吐量，而没有因此降低模型质量。
通过训练两个版本的GPT2模型（分别有83亿和25亿参数）来展示Optimus-CC的高效性。在我们配备了高端互联的大规模GPU集群环境中，我们在训练上获得了显著的加速。

背景阅读

3D并行

数据并行
管道并行
张量并行

这部分读者有兴趣可以自行搜索一下，csdn和知乎都有很多的公开文章介绍

Megatron-LM

Megatron-LM是一个使用3D并行训练极大规模NLP模型的框架。图1和图2显示了Megatron-LM的细节。它应用张量并行来分割模型层，如图2所示。为了减少通信时间，每个张量并行组都被放置在一个服务器节点内，这样它的通信就可以利用高带宽的服务器内部连接（即，NVLink）。多个服务器节点被用于数据并行和流水线并行。

Megatron-LM还使用交错调度和1F1B调度来减少管道气泡和峰值内存使用。

由于数据并行和流水线并行的通信在节点间网络中进行，这两种通信经常成为Megatron-LM的3D并行中的瓶颈。

本文工作的目标是减少这种通信的量，从而最大化大型NLP模型训练的吞吐量。

图一：
在这里插入图片描述
图二：

梯度压缩（这块是我比较想了解的地方）

这里保留了论文的引用，为了之后找相关论文的时候使用，我最近在看压缩相关的东西

梯度压缩常用于减轻数据并行中巨大的参数梯度通信成本。Topk、量化和低秩近似是梯度压缩的三种常见方法。

Top-k

基于Top-k[44, 64]的方法选择每层梯度的前k个元素进行压缩。在这种情况下，Top-k选择涉及排序开销，这在快速训练中非常关键。一些工作[11, 12]使用准排序来减少这种开销。Top-k方法的一个问题是，随着数据并行中GPU数量的增加，需要通信的元素总数会线性增加，因为每个GPU都会独立选择自己的k个元素。此外，Top-k方法需要一个额外的gather操作来选择索引，这在多工作器环境中比实际的参数梯度共享带来更多的开销。ScaleCom[12]通过使用每个工作器中梯度的Top-k索引相似性，解决了由于gather操作导致的这个梯度积累问题。Ok-Topk[41]也使用其高效的Top-k阈值估计算法成功解决了这些问题。
基于量化的方法将梯度量化以减少通信。TernGrad[81]使用三元值（-1,0,1）来极大地减少通信。AdaComp[11]还结合了残差加和和量化，以最小化损失压缩的误差。SignSGD[5]使用动量的符号来大幅量化梯度。1-bit Adam[70]通过量化使用warm-up迭代后优化器变量的稳定性，减少了Adam优化器的通信。
低秩近似使用矩阵分解来减少梯度的总通信成本。分解成本是矩阵分解的主要瓶颈。PowerSGD[75]通过只进行一次power迭代（对于经典的SVD是必需的）来减少这个成本。它重用前一次梯度压缩阶段的分解矩阵，以最小化压缩的误差。

这些方法的一个共同缺点是，所有的Top-k、量化和低秩近似方法都是固有的有损的，可能会降低模型的准确性。许多方法试图通过考虑动量[44]，提供误差反馈[11, 75]，或估计梯度运动[35, 83]来缓解这个问题。Optimus-CC也不例外，因为它采用低秩近似来压缩通信。我们提出了几种方法来在压缩3D并行的通信的同时，保持模型的质量。

重点关注这个低秩近似的方法，这会另开一篇博客分享阅读笔记，本文用了这个方法

本文的动机

在这里插入图片描述
注意看，这个通信时间占比很高，并且naive compress会使得validation 的PPL（困惑度）提升

本文的优势就出来了，ppl不变，通信时间减少。

讲一个故事理解本文的动机吧：

想象一个巨大的城市（NLP模型），这个城市由128个区域（GPU）组成，这些区域分布在16个大岛屿（计算节点）上。这些岛屿通过一系列高速公路（200Gbps Infiniband HDR互连）互相连接。城市的运作需要各区域之间的信息交流，这些信息交流可能是数据并行通信（DP Comm.），管道阶段之间的通信（Inter-stage Comm.），或嵌入同步（EMB Comm.）。

在每个岛屿上，区域之间的交流通过一个更快的交通系统（600GBps NVLink）进行，这个系统由于其高速度，使得岛内的通讯几乎不耗费时间。然而，岛屿间的通信却变得非常耗时，即使使用了高速公路。

为了提高城市运作的效率，**他们尝试了一种简化的方法，即通过压缩数据来减少岛屿之间的通信量（‘naive DP’和’naive CB’）。**这就像是将货物压缩成更小的包裹以便于运输。然而，这种方法导致了一个问题：信息在压缩和解压缩过程中丢失，使得城市运作的质量下降。

于是，一种新的方法Optimus-CC被提出。这种方法精心设计了压缩算法，使得岛屿间的通信只占据可忽略的时间，同时保持城市运作的质量不变，就像是找到了一种无损压缩的方法。这种新的方法成功地将城市运作的时间从8天减少到了6.97天，而且保持了城市运作的质量。

然而，当他们尝试将这种压缩方法应用到区域间的通信上（Opt-CC TopK），结果并不理想，因为这种通信方式并不适合压缩。这就像是一种压缩方式可能适合长途运输的大包裹，但可能不适合短途运输的小包裹。

所以这个方法只是局部的适用，用的地方要选对

本文的方法

在这里插入图片描述
图4a展示了每个管道阶段的基本1F1B调度，其中包括通信考虑因素。蓝色框代表前向传递，蓝色箭头表示阶段间前向通信。绿色框指的是反向传递，每次反向传递的时间大约是前向传递的两倍。绿色箭头代表阶段间反向通信。在反向传递完成后，发生数据并行通信（DP）。这种通信涉及到其他数据并行中同一阶段的GPU，如图4中的黄色框所示。此过程还包括了嵌入层参数的通信（EMB DP）。

本文方法旨在减少三个特定区域的通信量，如图4b所示。

OptimusCC方法旨在通过三种主要技术减少通信量：

压缩反向传播

此技术压缩阶段间通信，在图中将梯形的绿色区域变形为更接近矩形的形状，以缩短执行时间。为了防止模型质量下降，引入了两种技术，lazy错误传播和尾声压缩。

压缩目标：针对管道并行中的阶段间反向传播。
方法：利用懒散错误传播和仅在尾部进行压缩的策略，对阶段间的激活梯度进行压缩。

在进行压缩反向传播的过程中，我们采用低秩近似来减少阶段间反向传播的流量（激活梯度，如图4a和4b的绿色箭头所示）。由于这种通信是在每个阶段的计算之间发生的，因此它对性能的影响尤其明显，特别是在有大量管道阶段的情况下。理论上，**压缩前向流量也能提供类似的加速效果，但这可能会严重损害模型的收敛性。**虽然简单地压缩反向流量也会破坏模型的收敛性（如图3所示），但我们提出了两种策略——懒散错误传播和仅在尾部进行压缩，以保持模型的收敛性。

懒惰错误传播 lazy error Propagation.

这是一种新颖的技术，它能够实现压缩反向传播而不降低模型质量。

简单地将压缩应用于阶段间反向传播通信会导致模型质量严重下降，这是因为压缩算法的有损性质导致的错误。
在这里插入图片描述

图5展示了带有懒散错误传播的阶段间反向传播。当使用懒散错误传播时，压缩后的数据发送到早期阶段进行反向传播后，错误将被保留在内存中。这个保留的错误会被添加到下一个微批次的反向流量中。

例如，在第三个微批次中，设备2生成了来自样本16-23的梯度∇𝑓 16−23 (𝑌1)。梯度被压缩并发送到设备1，同时压缩错误𝜖16−23 1被保留在内存中。在下一个微批次中，当生成∇𝑓 24−31 (𝑌1)时，𝜖16−23 被添加进来。这个和在被发送前被压缩，新的错误𝜖24−31 被保留。

也就是！把压缩和错误解耦了！压缩先行，错误后行，错误将给延迟更新！

尽管错误稍微延迟到下一个微批次，但对模型质量的影响几乎可以忽略不计，因为模型的更新只在所有微批次处理完之后才发生。换句话说，懒散错误传播不会受到权重陈旧性效应的影响，因为所有的微批次仍然是基于权重的同一版本来执行的。

懒散错误传播将特定微批次的错误延迟到后续微批次的错误中。为了让懒散错误传播起作用，结果的平均梯度应该正确地近似于不使用压缩时得到的平均值。一种直观的理解是，梯度是在整个小批次上累积的，而不像激活那样直接影响模型的结果。正因为如此，我们认为在后续微批次中传播的错误不会大大影响总和。我们发现，当错误的分布与激活的分布无关时，这一点是正确的。以下，我们会提供对这个条件的数学分析。

这可能就是为什么用在反向传播，而不是前向传播的原因了。前向的时候要是错误在下一个时间到来，这一个时间的模型就算不对了。

论文中还有一段数学证明，这段非常的数学，如果你想知道他具体是如何证明这个压缩后恢复是无损的，建议你阅读原文。这个数学推导不方便打字。

仅尾声压缩 (Epilogue-Only Compression)

虽然压缩技术显著减少了通信时间，但过度压缩总是有降低模型质量的风险，这本质上成为了训练速度和模型质量之间的权衡。

令人惊讶的是，在管线并行（pipeline parallelism）中，我们可以选择压缩位于关键路径上的数据。在图4a中显示的基线调度下，很多来自阶段间反向传播的通信几乎完全与计算重叠（隐藏）。然而，位于关键路径上的阶段（称为尾声）的通信并未被计算隐藏，如图6a所示。
在这里插入图片描述

融合嵌入同步

此技术针对共享嵌入层，这些层生成两个all-reduce通信。该方法将这两个all-reduce通信融合为一个all-reduce通信，减少了流量。

当模型中的任何权重参数在网络中多次使用时，这些参数会从多条路径上累积梯度。如果这些路径都在单个GPU内部，那么这不会引发任何问题。然而，如果它们在不同的GPU上执行（无论是使用张量并行还是管道并行），它们就需要对参数梯度进行同步，这是另一种类型的通信。

在我们主要关注的大型自然语言处理模型中，这样的结构通常会出现在嵌入层，如图7中紫色框所示。
如果网络的输出采用文本格式（如在预训练阶段），嵌入层会被使用两次：一次在输入处，将单词转换为嵌入值，另一次在将输出向量转换为单词。因为它们位于模型的开始和结束位置，所以如果使用管道并行，它们总是会生成一个节点间通信。

在传统的策略中，我们会在两个阶段进行这样的同步：

数据并行阶段：在这个阶段，我们将训练数据集划分为多个子集，并在每个计算节点上对一个子集进行处理。每个节点会独立计算其数据子集的梯度。
all-reduce阶段：在这个阶段，我们将所有计算节点上计算出的梯度进行合并，并将结果广播回所有节点。这样，每个节点都可以使用同步后的梯度更新其模型参数。

"融合嵌入同步"的策略是将这两个阶段的通信融合为一个。这是通过将与嵌入层相关的两个all-reduce通信（一个来自数据并行方式，另一个来自嵌入同步）融合为一个all-reduce来实现的。这样，我们就将两个通信阶段减少为一个，从而有效地减少了通信的总成本。
在这里插入图片描述

选择性阶段压缩

此技术针对数据并行流量，但仅压缩那些在关键路径上的元素。这种方法比传统的数据并行流量压缩提供了更好的模型质量和执行速度之间的平衡。

当我们同时使用数据并行和管道并行时，会出现一个问题：在数据并行阶段，我们需要对所有节点上的梯度进行同步，这会消耗大量的通信带宽。为了解决这个问题，我们可以使用压缩技术来减少需要传输的数据量。但是，这种压缩方法可能会导致模型质量的下降，因为压缩操作可能会丢失一些重要的信息。

"选择性阶段压缩"策略的提出，就是为了解决这个问题。具体来说，我们不是简单地对所有阶段的数据进行压缩，而是选择性地对某些阶段的数据进行压缩。这些被选择的阶段通常是那些在管道并行中提前完成的阶段。因为这些阶段可以提前开始通信，所以我们可以将他们的数据先进行压缩，然后再进行传输。这样，我们既可以减少通信带宽的消耗，又可以尽可能地保持模型的质量。

在这里插入图片描述