大规模分布式 AI 模型训练系列——序列并行

程序员笑武

于 2024-09-28 20:31:27 发布

阅读量970

点赞数 17

文章标签：人工智能计算机视觉深度学习 neo4j 大数据

本文链接：https://blog.csdn.net/m0_59164304/article/details/142621197

版权

一、背景

之前的文章中我们详细介绍了大规模分布式训练中的数据并行（Data Parallelism，DP）、张量并行（Tensor Parallelism，TP）、流水线并行（Pipeline Parallelism，PP）以及专家并行（Expert Parallelism，EP）。这篇文章中我们继续介绍 LLM/LMM 长序列场景经常使用的序列并行（Sequence Parallelism），以及相关的各种优化手段。

二、新加坡国立 Sequence Parallelism

2.1 摘要

在 Transformer 模型中，自注意力机制的内存需求、计算量与序列长度成二次方关系，导致序列比较长时可能存在明显瓶颈。在 [2105.13120] Sequence Parallelism: Long Sequence Training from System Perspective 中，作者提出了序列并行（Sequence Parallelism），这是一种内存高效的并行方案，可以帮助突破输入序列长度限制，并在 GPU 上高效地训练更长的序列。提出的方法与大多数并行方案兼容，比如 DP、PP 和 TP，也就是说可以使用 4D 并行。

此外，不再需要在单个设备上保留整个序列，具体来说，作者将输入序列分成多个 Chunk，并将每个 Chunk 输入到相应的设备。为了计算注意力输出，将环形通信与自注意力计算相结合，提出了 Ring-Self-Attention（RSA）。实验表明，与 TP 相比，扩展到 64 个 NVIDIA P100 GPU 时，提出的方法可以实现 13.7x 的最大 Batch Size，3x 的序列长度。

PS：本文的方案还是针对 Bert 这种 Encoder Only 的模型。针对 Decoder Only 的 GPT 模型，由于 Attention Score 是个下三角矩阵，需要进一步考虑负载均衡问题。

2.2 方案概览

如下图 Figure 1 所示为本文的 Sequence Parallelism 方案与 PP 和 TP 的区别，其核心是每个设备分到一部分 Token，而且每一层也都需要通信：

2.3 Ring Self-Attention

如下图所示为标准的 Self-Attention 计算方式，其核心为两个矩阵乘法：Q * KT 和 Attention Score * V：

如下图 Figure 2 所示为本文提出的 Ring Self-Attention，作者将其按照上述两个矩阵乘法分为两个 Step，每个 Step 实质上都是矩阵乘法的分块计算。假设序列长度为 L，设备数为 N：

Q * KT：每个设备上 L/N 个 Token 的 Query 和 Key，通过 Ring 的方式传递 Key。N-1 次传输，N 次计算后得到完整的 Attention Score，每个设备上只有 1/N 的 Attention Score。
Attention Score * V：每个设备上 L/N 个 Token 的 Value，通过 Ring 的方式传递 Value。N-1 次传输，N 次计算后得到完整的 Attention Output，每个设备上只有 1/N 的 Attention Output。

如下图所示为一个完整的计算过程，第一行为 Step 1，第二行为 Step 2。核心就是分块矩阵乘，每次计算一个 Block，通过移动 Key 和 Value 可以实现计算不同的 Block：

2.4 FFN 切分

我们在之前的文章中提到过，NVIDIA 在 [1909.08053] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism 中发布了 Megatron-LM 框架，其中使用了 Tensor Parallelism 的方案。如下图所示，其 MLP 层的两个 FC 采用先列切（A，Column Parallelism），然后行切（B，Row Parallelism），这样两个 FC 之间不用通信。同时每个设备上都只有权重参数 A 和 B 的一部分，也可以减少显存的占用：

在上述方案中，每个设备上都有完整的 FFN 输入 X 和输出 Z，当序列比较长时，输入和输出占据的显存空间反而可能比权重 A 和 B 更大。此外，FFN 中 Token 之间是没有交叉的，也就是每个 Token 可以单独计算，基于这个特性，作者提出了 FFN 的序列并行方案（每个设备上需要存储完整的 A 和 B）。如下图 Table 1 所示为 Tensor Parallelism 和 Sequence Parallelism 的显存占用：

可以看出，当满足以下条件，也就是 BL > 32H 时，序列并行会更加节约显存：

2.5 结果

如下图 Figure 5 所示，随着 Parallel Size 的增加，Sequence Parallelism 会更有优势。（PS：由于当前硬件配置，通常单机内会有 8 个 GPU，通过 NVLink + NVSwitch 高速互联，通信带宽很高，而跨机后只能通过网络互联，通信带宽低得多，因此 TP 的 parallel Size 往往不会超过 8）

三、NVIDIA Megatron-SP

3.1 摘要

NVIDIA 在 [2205.05198] Reducing Activation Recomputation in Large Transformer Models 中提出了针对 Transformer 模型的选择性激活重计算（Activation Recomputation）以及序列并行（Sequence Parallelism）方案来减少内存开销。结合 Tensor Parallelism，几乎可以消除重新计算激活的需求。

基于提出的方案，可以在 1T 参数量的语言模型训练中将激活内存减少 5x，同时将激活重计算的执行开销减少 90% 以上。比如，在 2240 个 NVIDIA A100 GPU 上训练 530B 参数的 GPT-3 模型时，实现了 54.2% 的 MFU，比使用重计算的方案（42.1%）快 29%，相应实现都集成到了 Megatron-LM 中。

PS：需要说明的是，Megatron-SP 需要配合 Megatron-TP 运行。

3.2 方案

如下图 Figure 4 所示为 Transformer Layer 中标准的 Tensor Parallelism 分布式方案，其核心是两个（Attention 和 FFN）连续矩阵乘的切分方案（先列切、然后行切）：

使用上述切分方案，每个设备每层的 Activation 占据的显存空间如下所示，其中 s 表示序列长度，b 表示 Batch Size，h 表示 Hidden Size，t 表示设备数，a 表示 Head 个数（假设激活都占用 2 个字节，Mask 占用 1 个字节）：

由于 LayerNorm 和 Dropout 的计算量很小，因此这一部分并没有切分，在每个设备上都会独立、重复的计算，所以对应的 Activation 也并没有切分。上述公式中的 10sbh 对应着未切分的 Activation，包括 LayerNorm 和 Dropout 的输入各 2sbh（各 2 层，共 8sbh），以及 2 个 Dropout 的 Mask 各 sbh，共 10sbh。然而，实际上这一部分 Token 之间并没有交叉，可以按 Token 分到不同的 Device 独立计算，作者称这种方式为 Sequence Parallelism，如下图所示，对应的通信操作也从 f 变为 g。

按上述 Sequence Parallelism 方式切分后，10sbh 的 Activation 也可以明显降低，每个设备只需要 1/t，最终每个设备每层需要的 Activation 如下所示：

3.3 结果

如下图 Table 2 所示，作者总结了各种方案每个设备每层的 Activation 大小，通过 TP + SP 可以将 Activation 降低到 1/t：

如下图 Figure 7 所示，作者验证了不同规模模型在引入 Sequence Parallelism 和选择性激活重计算方案后的激活量。基线为 Tensor Parallelism，可以看出，仅引入 Sequence Parallelism 就可以将 Activation 降低到 60% 左右：

四、伯克利 Blockwise-Parallel-Transformer

4.1 摘要

在 [2305.19370] Blockwise Parallel Transformer for Large Context Models 中，为了解决长序列场景中显存的挑战，作者提出了 Blockwise Parallel Transformer（BPT），利用 Blockwise 计算以及 Self-Attention 和 FFN 融合来最大限度降低内存开销，以便在保持内存效率的同时处理更长的输入序列。使用 BPT，可以训练比原生 Transformer 长 32x 的序列，甚至也比 FlashAttention 等内存高效方案支持的序列长度长 4x。

PS：其中的所有对比实验都使用了激活重计算（Activation Recomputation），如果不使用则无法实现节约这么多的显存，使用和不使用激活重计算的吞吐并没有对比。

4.2 方案

本文的核心思路是利用了 Decoder Only 的 Transformer 模型的 Causal 特性，任何一个 Token 都只依赖之前的 Token，极端情况下甚至可以逐个 Token 计算，而反向传播需要的中间 Activation 可以通过激活重计算的方式节约。

如下图 Algorithm 1 所示，其核心思路是将输入 Sequence 拆分成不同的 Block 依次计算，其包含两层循环：

外层循环：按照 Query 切分，计算当前 Query Block 对应的完整 Attention 和 FFN。
内层循环：一个 Query Block 内部的训练，分为如下几个 Step
遍历 Key 和 Value 的 Block，依次计算 Attention Output
对 Attention Output 进行修正（PS：和 FlashAttention 一样）
计算 FFN 以及残差连接

4.3 结果

如下图 Table 2 所示，使用本文的方案可以大幅提升可支持的序列长度：

由于无论如何分块，整个序列的中间激活在反向传播中都是需要的，只能通过时间换空间（激活重计算）的方式节约，如下图所示，甚至整个 FFN 的中间激活都会丢弃，只保留当前 Block，这也是为什么能比 FlashAttention 省显存的主要原因，而这里也会极大增加反向传播的计算量：

五、微软 DeepSpeed-Ulysses

5.1 摘要

在 [2309.14509] DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models 中，作者指出之前的 Sequence Parallelism 会受到内存通信效率低下的制约，限制了长序列大模型的可扩展性。因此，作者提出了 Deepspeed-Ulysses，用于实现具有极长序列的高效、可扩展 LLM 训练。

DeepSpeed-Ulysses 的核心是沿序列维度对输入数据进行切分，并采用高效的 All-to-All 集合通信进行注意力计算。理论通信分析表明，其他方法会随着序列长度增加而产生通信开销，而 DeepSpeed-Ulysses 在序列长度和计算设备按比例增加时保持恒定的通信量。评估实验表明，与当初 SOTA 基线相比，DeepSpeed-Ulysses 的训练速度提高了 2.5x，序列长度增加 4x。

PS：为了充分节约显存，这里的 DeepSpeed-Ulysses 会与 DeepSpeed-Zero3 结合使用，并且需要 Attention head 个数是 SP 设备数的整数倍。

5.2 方案

如下图 Figure 2 所示为 DeepSpeed-Ulysses 的核心方案，其思路很简单：

假设有 P 个设备，则将输入序列 N 均分到 P 个设备，每个设备 N/P，除了 Attention 部分之外，其他部分都可以独立计算。
在 Attention 部分则按照 Head 切分，假设有 hc 个 Head，则每个设备负责 hc/P 个 Head 的计算，也就是 Head 个数需要是设备数的整数倍。
在 Attention 计算之前，通过 All-to-All 操作将 Head 对应的 Q、K、V 分发到对应设备，通信量为 3Nd；计算完之后再通过 All-to-All 将 Attention Output 发送回对应的设备，通信量为 Nd。

如下图所示为上述 Q、K、V 的切分方式及 All-to-All 通信传输（如果 Q、K 和 V 分开存储，需要 3 次独立的 All-to-All，但实际上可以将 Q、K 和 V 拼接存储，这样只需 1 次 All-to-All），其中 T3-H2 表示 Token 3 对应 Head 2 的输入。完成 All-to-All 之后每个 Head 都可以独立的计算，计算完之后将 Attention Output 按照反向的 All-to-All 重新传输即可。

5.3 方案对比

如下图 Table 1 所示，作者与我们之前介绍的“新加坡国立 Sequence Parallelism”（ColAI-SP）和 Megatron-SP 进行了对比，可以看出，它们都可以有效节约显存空间。

ColAI-SP：每个设备都要见过所有的 K 和 V，假设消息大小为 M，则每个设备的通信量为 O(M)；此外，并没有切分参数；主要针对 Encoder 的 Attention。
Megatron-SP：每一层都需要两个 AllGather 和两个 ReduceScatter，因此每个设备的通信量为 O(M)；此外，需要结合 DeepSpeed-TP 使用。
DS-Ulysses：All-to-All 总的通信量是 M，因此每个设备是 O(M/P)；此外，借助 DeepSpeed-Zero3 实现了参数的切分。

5.4 结果

如下图 Figure 4 所示，作者在 32 个 GPU 上对比了 DeepSpeed-Ulysses 和 Megatron-LM 训练 7B LM 的 TFLOPs，可以看出，DeepSpeed-Ulysses 可以获得更高的性能，并且可以支持更长的序列长度。

六、伯克利 Ring-Attention

6.1 摘要

Blockwise-Parallel-Transformer 的作者基于之前的工作又提出了 Ring Attention 方案（[2310.01889] Ring Attention with Blockwise Transformers for Near-Infinite Context）。其 Attention 计算可以看做一个分布式的 FlashAttention，此外，作者将 Key、Value 的通信与分块注意力的计算完全 Overlap。提出的方案可以训练比之前方案长几倍的序列，而无需近似或者考虑额外的通信和计算开销。使用提出的方案，作者在语言建模任务上实现了数百万 Token 的长上下文训练。

PS：Ring-Attention 依然需要和激活重计算（Activation Recomputation）结合使用，可以参考 https://github.com/forhaoliu/ringattention/issues/4。除此之外，本文的 Ring-Attention 其实和前面“新加坡国立 Sequence Parallelism”中的 Ring-Self-Attention 思路是一样的，只是结合了 FlashAttention 实现更加高效。我们在困惑 Ring-Attention 为什么没有引用 Ring-Self-Attention 的时候发现大佬也提出过相关的问题：

6.2 方案

作者之前提出的 Blockwise-Parallel-Transformer 还是一个 Local 的分块方案，虽然可以训练更长的序列，但是没有结合分布式训练，其训练效率会很低，此外依然需要保存每层的中间结果，大小为 2bsh，随着序列长度的增加会线性增加。

因此本文中作者将其扩展到了分布式方案，如下图所示，每个设备上一部分的 Sequence Block，然后同样以 Ring 的方式同时传输 Key 和 Value，同时计算相应的 Attention Score 和 Output。在计算的同时可以将 Key 和 Value 传输到下一个设备。

在这里插入图片描述

如下图所示，在我们之前介绍的 ColAI-SP 中，其是通过两次 Ring（一次 Key，一次 Value）来实现的，然而在我们观察第二行的结果时可以发现，实际上第二次 Ring 训练的第 k 次计算只用到第一次 Ring 循环第 k 次的输出（如果不考虑 Softmax），而 Softmax 带来的差异又可以通过 Online-Softmax 的方式修正（也就是 FlashAttention）。

基于以上的思路，实际上基于一次 Ring 就可以同时计算 Attention Score 和 Attention Output，只不过计算之后还需要额外的修正操作。如下图所示：

当然，上述的 Blockwise Attention 也会和 Blockwise FFN 相结合，以便充分降低显存的开销。同样的，在所有实验中作者都使用了 Gradient Checkpointing 方案。

6.3 结果

如下图 Table 1 所示为本文的 Ring-Attention 方案与其他方案的 Activation 占用对比，其中 c 表示每个设备上的 Block Size：

如下图所示，本文提出的 Ring-Attention 能支持训练的序列长度远大于之前的方案：

七、伯克利 DistFlashAttention

7.1 摘要

由于 Decoder 的 Causal 特性，Attention 是一个下三角形式，Ring-Attention 的切分方式会导致负载不均衡。[2310.03294] DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training，作者了更加均衡的调度方案。

7.2 方案

如下图 Figure 1 所示，左图为原始的 Ring-Attention 方案，可以看出，靠前的 Worker 会存在很多无效的等待，导致算力的浪费。右图所示为本文的方案，在前面 Worker 等待的时候可以令其执行本来属于后面 Worker 的 Workload，以此可以充分的减少 Bubble。图中的示例，优化前需要 8 个 Step，而优化后只需 5 个 Step 即可，大幅降低运行时间。

7.3 结果

如下图 Table 1 所示，与优化后的 DistFlashAttention 相比 Megatron-LM 在长序列训练时具有比较明显的优势：

八、腾讯 USP

8.1 摘要

我们在之前的部分介绍了 DeepSpeed-Ulysses（SP-Ulysses）和 Ring-Attention（SP-Ring），它们都在序列维度对输入 Tensor 进行切分，但是在 Attention 计算时采用了不同的切分方式来支持更长的序列长度。在 [2405.07719] USP: A Unified Sequence Parallelism Approach for Long Context Generative AI 中，作者对其进行统一，提出 USP（Unified Sequence Parallelism）。具体来说，该方法结合了 SP-Ulysses 和 SP-Ring 的优势，对 Transformer 模型结构和网络硬件拓扑更加鲁棒。使用提出的 USP，作者在两个 8xA800 节点上实现了 47% 的 MFU，可以以 208K 的序列长度训练 LLaMA3-8B 模型。

8.2 方案

我们之前详细介绍了 SP-Ulysses 和 SP-Ring 的方案，但是并没有介绍其局限性，这里简单总结一下：

SP-Ulysses：按照 Head 切分，Head 数是 SP 并行度的整数倍，也就意味着 Head 数目比较少时并行度不高，比如现在流行的 GQA，影响了其可扩展范围。
SP-Ring：通常来说包括 3 个方面
可以看成是一个分布式的 FlashAttention，采用 Ring-Attention 切分为很多小的矩阵乘，可能导致在 GPU 上的计算效率不高，通信的问题也会更加突出。
由于 Decoder 的 Causal 特性，存在负载不均衡问题。
这种 Ring 的方式可能受到网络拓扑的限制，比如机内 NVLink 高速互联，机间 IB 互联，则通信带宽会受到机间 IB 带宽的限制。

在这里插入图片描述

针对 SP-Ring 负载不均衡的问题，作者采用了另一种调度方式。如下图 Figure 3 所示，序列切分时采用对折的方式即可，可以保证每个设备的计算量近似相等（可以参考这个工作 [2311.09431] Striped Attention: Faster Ring Attention for Causal Transformers）：

关于其他几个问题，可以采用 SP-Ulysses 和 SP-Ring 相结合的方案解决。如下图 Figure 4 所示，假设有单机 8 个 PCIe GPU，每个 CPU NUMA 下有 4 个 GPU，也就是说，不同 CPU 下的 GPU 通信需要通过低带宽的 QPI/UPI，可能成为瓶颈。假设序列长度为 8K，有 4 个 Head，作者的方案可以看做一个 2D Attention：

在同一个 CPU 下通过 PCIe 互联的 GPU 使用 SP-Ulysses，这几个 GPU 通过 All-to-All 通信；共两组，每组负载 4K 序列长度。这一组内是按 Head 切分，每一个 GPU 处理一个 Head，也就是每个 GPU 处理 4K 序列长度，1 个 Head，矩阵乘维度更大，计算效率更高。
两个 CPU（两组 SP-Ulysses）之间采用 SP-Ring，可以将序列长度从单个 SP-Ulysses 的 4K 进一步扩展为 8K。此时只需要两个 CPU 之间的 Ring，避免了 8 个 GPU Ring 时 QPI/UPI 带宽成为瓶颈。除此之外，也克服了 SP-Ulysses 方案 4 个 Head 最大只能使用 4 个 GPU 的并行度约束。

当然，其实更常见的场景是单机内 8 个 GPU 通过 NVLink+NVSwitch 高速互联，而机间通过 IB/RoCE 互联。以 LLaMA3-8B 为例，其采用了 GQA，只有 8 个 Attention Head，因此可以按照下述的方案：

在机内使用 SP-Ulysses，每个 GPU 负责一个 Head，也就是 SP-Ulysses 的 degree 为 8；SP-Ulysses 的 All-to-All 可以使用高速的 NVLink。
机间使用 SP-Ring，每台机器负责一个 Block，假设 4 台机器，也就是 SP-Ring 的 degree 为 4；SP-Ring 的 P2P 通信可以使用 IB/RoCE。

8.3 分析

如下图 Table 2 所示，作者基于 GPT Style 的模型（比如，FFN 2 层 Linear，升维维度为 4N；没有考虑 GQA）分析了不同分布式方案的通信开销和显存占用大小：

基于分析作者也提供了几个小的 Tips：

尽可能优先使用 DP 而不是 SP。只有当 Batch Size 不足以更好的切分时考虑是否使用 SP。
在使用 SP 时应始终与 ZeRO-1/2 结合使用，也可以进一步考虑采用 ZeRO-3 和 Offload 方案来平衡通信开销以及节约的内存。
当涉及大规模分布式时 SP 的通信成本优于 TP-SP，GQA 可以进一步降低 SP 的通信成本。
将 TP-SP 切换为 SP 并不会增加可支持训练的序列长度，SP+Zero3 可以训练与 TP-SP 相似的序列长度。
当 Head 数目有限时，可能需要设置一个比较大的 SP-Ring degree，以便在更多的设备上训练长序列，这是 TP-SP 无法实现的优势。
在 4D 混合并行中，Process Group 的维度从低到高的顺序为 TP、SP-Ulysses、SP-Ring、ZeRO-DP、PP。

8.4 结果

如下图 Table 5 所示，在单个 8xA800 NVLink 节点内训练 LLaMA2-7B 模型，最优方案都没有使用 SP-Ring：

在这里插入图片描述

如下图 Table 6 所示，在两个 8xA800 NVLink 节点内训练 LLaMA3-8B 模型，当固定 Global Batch Size 为 1 时往往需要结合 SP-Ring，但在这个规模往往不需要太大的 SP-Ring degree：

九、上海 AI Lab&北大 LoongTrain

LoongTrain 和腾讯 USP 是同期的工作，方案基本差不多，核心也是 2D-Attention 来结合 DeepSpeed-Ulysses 和 Ring-Attention，也解决了 Ring-Attention 的负载均衡问题。同样基于 ZeRO 实现，并且与 Selective Checkpoint++ 相结合。实验结果表明，LoongTrain 在端到端训练速度和可扩展性方面都优于最先进的基线，即 DeepSpeed-Ulysses 和 Megatron Context Parallelism，并将模型 FLOPs 利用率（MFU）提高了 2.88×。

对应的论文为：[2406.18485] LoongTrain: Efficient Training of Long-Sequence LLMs with Head-Context Parallelism

PS：由于其思路类似，所以这里我们不再展开，不过其中的很多分析和实践还是很值得读一下的。

十、NVIDIA LongVILA

在 [2408.10188] LongVILA: Scaling Long-Context Visual Language Models for Long Videos 中，为了支持长视频（序列很长）场景，作者同样采用了类似 USP 和 LoongTrain 的方案，如下图 Figure 5© 所示为其核心的 2D Attention，主要是将其扩展到了多模态场景。

作者也同样讨论了 SP-Ring 在异构网络下的通信瓶颈问题，也更进一步讨论了多模态场景下 Token 切分导致的负载不均衡问题（PS：核心是说之前的方案会按图片维度切分，导致最后一个 Worker 还要负责 Prompt 相关的文本 Token，导致最后一个 Worker 要处理的 Token 更多）。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述