[张量/序列并行]图解 DeepSpeed-Ulysses & Megatron-LM TP/SP

来源丨https://zhuanlan.zhihu.com/p/...

0x00 前言

最近正在为打工人的内卷热潮，略尽自己的一份绵薄之力，已经挺长一段时间没有更新了。正好工作中需要用到各种序列并行的技术，于是打算开几篇小文章来记录一下自己的理解。对比起来 DeepSpeed-Ulysses Sequence Parallelism，并不是十分好理解，还是 Megatron-LM 的 Tensor/Sequence Parallelism 更直观一些。因此本文将通过图解的方式，由浅入深记录一下自己对 DeepSpeed-Ulysses 和 Megatron-LM TP/SP 的理解。

本人更多的技术笔记以及 CUDA 学习笔记，欢迎来 CUDA-Learn-Notes（CUDA Learn Notes with PyTorch）查阅。CUDA-Learn-Notes 包括了本人的 LLM/VLM 文章整理，以及对于 SGEMM/HGEMM/GEMV 等常见 CUDA Kernel 的示例实现，目前已经累计 2k+ stars，传送门：https://github.com/xlite-dev/...

CUDA Learn Notes with PyTorch

0x01 Megatron-LM Tensor Parallelism

首先，简单讲一下 Megatron-LM 的 Tensor Parallelism 和 Sequence Parallelism[1]，我们知道，Megatron-LM 在 2020 年的论文中，针对 Transformers 架构，最开始提出的并行化方式是 Tensor Parallelism，也就是以下这样经典的图：

Megatron Tensor Parallelism

Megatron Tensor Parallelism 对 Transformers 结构中的 Self Attention 和 MLP(Linear)模块进行张量并行，从而可以这两个模块的权重和计算都平均地分摊到 P 个 GPU 上。对 MLP 和 Attention 采用的张量并行策略，具体如下：

MLP Tensor Parallelism

因为矩阵乘是可以分块计算的，从而使得对 MLP 实现 Tensor Parallelism 成为可能。

矩阵乘分块

上一小节提到，Tensor Parallelism 的基础是矩阵乘可以分块。我们知道，在 CUDA 中，通用矩阵乘 GEMM（SGEMM/HGEMM 等），在工程实现上，是通过分块矩阵乘，利用高带宽的 SRAM 以及片上存储-寄存器（数据复用）等技术方式来实现的。把分块矩阵乘这个概念从单卡扩展到多卡，就相当于每个 rank 处理一个矩阵的其中一个分块，而整个矩阵乘的结果则通过 NCCL 等集合通信库提供的接口进行聚合。

Parallel GEMMs

从权重矩阵 A 的视角去看，矩阵 A 分块的切分方法有两种，分别是按【行】切和按【列】切。

权重 A 行/列切分

Row Parallel

假设现在有 2 卡，如果对权重矩阵 A 执行按【行】切分的方式，得到 A1, A2 分块的维度为[d/2,d]，那么为了矩阵乘可以顺利进行，需要将 X[N,d]按照列方向进行切分，即得到 X1, X2 分块的维度为[N,d/2]，分个分块各自计算自己的结果[N,d]，最后需要通过 All-Reduce 将结果累加，得到最终的输出：Y=Y1+Y2=X1A1+X2A2。

Column Parallel

如果是对重矩阵 A 执行按【列】切分的方式，得到 A1, A2 分块的维度为[d,d/2]，那么为了矩阵乘可以顺利进行，X[N,d]需要保持不变，即每个 rank 上需要有一份完整的 X，分个分块各自计算自己的结果[N,d/2]，最后需要通过 All-Gather 将结果聚合，得到最终的输出：Y=[Y1, Y2]=[XA1, XA2]。

MLP 张量并行

既然存在权重矩阵不同的切分方式，那么对于 MLP 模块，我们的目标当然是希望找到一种方式，使得其通信量最低。首先，我们看一下 MLP 模块的计算公式：

可以发现，MLP 中存在两个矩阵乘以及 GELU 和 Dropout。GELU 和 Dropout 是逐元素操作，怎么切分对他们都没影响，可以忽略。而”(XA)B“，我们希望，为了通信量尽可能低，在 XA 后得到结果，可以直接用来和 B 进行矩阵乘，这是可以实现的。在 Megatron 中 MLP 权重矩阵 A、B 的切分方式为：对 A 进行列切分，B 进行行切分，那么有：（行列切分的概念，大家可以搜索其他 Tensor Parallelism 的文章进行阅读，推荐: 猛猿：图解大模型训练之：张量模型并行(TP)，Megatron-LM）

此时，可以让卡 1 计算图片，卡 2 计算图片，XA1 的结果可以直接在当前 rank 继续和 B1 相乘，不需要额外的通信，只需要在全部计算完成后进行一次 all-reduce 通信即可，这样，就可以只进行一次通信来完成计算。

Attention 张量并行

而 Attention 部分，由于 MHA 的每个 Head 都是独立不依赖的，因此天然就可以按照 Head 进行切分，把不同的 Head 放在不同的卡上进行计算即可，如下图所示。

Self Attention Tensor Parallelism

0x02 Megatron-LM Sequence Parallelism

未被分摊的激活值显存

回头看 Megatron 的 Tensor Parallelism，我们会发现 Tensor Parallelism 只处理了 Transformers 中 Attention 和 MLP 模块，LayerNorm 以及 Dropout 是保留原样的。也就是说，这两个算子，在每张卡上，对应的输入和输出都是完整的激活值。对于 long context 场景（比如 128K）来说，LayerNorm 以及 Dropout 计算占比不大，但是激活值显存占比却不小。张量并行后 per layer 的激活值显存占用计算公式如下：

张量并行后 per layer 的激活值显存占用

其中，加项 10sbh（s 表示 seqlen，b 表示 batch size，h 表示 hidden_size）来自与 LayerNorm，这部分是无法被张量并行分摊的。当序列长度越来越大，10sbh 会随着线性增加，从而成为训练和推理过程中的显存瓶颈。

序列并行+张量并行

那么，既然如此，一种很自然的方式就是对 LayerNorm 和 Dropout 在 seqlen 维度做并行，将显存压力分摊到所有使用的 GPU 上，因此，这种并行方式也叫做 Sequence Parallelism。

Megatron-LM Sequence Parallelism

我们可以看到，Megatron-LM 的做法是在 Tensor Parallel 前后，插入了 Sequence Parallel。这也是需要注意的点，这意味着，Megatron-LM 中的 Sequence Parallel 是要和它的 Tensor Parallel 一起使用的。而 DeepSpeed-Ulysses 的序列并行，则不需要和 Tensor Parallel 一起使用。Megatron-LM 中，张量并行的部分保持不变，只是对 Attention 和 MLP 的输入输出做了序列并行。

通信方式的改变

并行方式的变化意味着通信方式也需要跟着改变，原来纯 Tensor Parallel 需要 2 次 all reduce，而在 Sequence Parallel+Tensor Parallel 的情形下，则变成了 2 次 all gather 和 2 次 reduce scatter，根据论文的说明，这两种情况下的通信量是一致的。

激活值显存占用分析

在使用序列并行后，每 layer 的激活值需要的显存占用为：

SP: 每 layer 的激活值需要的显存占用

可以看到，10sbh 这个加项已经可以被分母 t 进行分摊了。最后，贴一下论文中给出的不同并行情况下 per layer 需要的显存（如果加上 recompute，显存需求会进一步下降，具体取决于配置）。

不同情况下 per layer 需要的显存

更详细的分析，推荐阅读：猛猿：图解大模型训练系列：序列并行 1，Megatron SP(https://zhuanlan.zhihu.com/p/...)

0x03 DeepSpeed-Ulysses Sequence Parallelism

DeepSpeed 的 Ulysses 是实现 Sequence Parallelism 的另一种方式，这种方式可以单独使用，不像 Megatron-LM 的 Sequence Parallelism，需要结合 Tensor Parallelism 一起使用。Ulysses 认为当前各种的序列并行的算法受限于他们次优的通信效率，无法在更长的序列上进行扩展（over a million tokens）。因此低通信量是 Ulysses 的核心卖点之一。

核心卖点：低通信量

主要创新

不过，刚开始看 DeepSpeed-Ulysses 的时候，着实是有点晕。特别是论文里边那张朴素的流程图，让我感觉真是听君一席话，如听一席话。好了，开始进入正题。首先，我们来看下朴素的 MHA 的做法。

朴素 MHA

在没有 Sequence Parallelism 的情况下，完整的输入张量[N,d]会被 broadcast 到各个卡上，然后每个卡通过完整的输入计算各自的 QKV，以及后续的 Attention 及 MLP。

DeepSpeed-Ulysses 整体流程

而 DeepSpeed-Ulysses 则是先对序列进行切分，然后再发送给不同的卡进行运算，具体步骤为，首先，沿序列维度切分 Batch 内的各个样本；然后，在计算注意力之前，使用 All-To-All 通信把 Query、Key 和 Value 进行聚合，以便每张卡上都具有完整序列长度，同时使得各张卡上只处理部分注意力头，以便并行计算注意力得分；最后，再使用 All-To-All 来沿着注意力头收集结果，同时沿着序列维度重新分区。

DeepSpeed-Ulysses Sequence Parallelism

然而，我并不认为大家在看完这段文字和图示后，能够很好地理解 DeepSpeed-Ulysses，特别是其中两个 All-To-All 所扮演的角色和作用。我自己刚开始也是如此，于是决定自己动手画一画加深理解，以下，是我根据自己的理解画的 DeepSpeed-Ulysses 序列并行流程图，如有错误，欢迎指正。接下来，我们就根据这张图流程图来简单说明一下 DeepSpeed-Ulysses 算法。

DeepSpeed-Ulysses 流程图

基本假设与符号说明

为简化理解，我们先假设使用的卡数 P 刚好等于 MHA 中注意力头数 head num。事实上，P 和 head num 应该需要满足整除关系，即 head num 可以被 P 整除。图中，Local O/Q/K/V 表示每个 rank 上的数据，Qh, Kh, Vh, Oh 表示当前 head 上在 N 维度上完整的数据。N 表示序列长度 seq len，d 表示 hidden size，head size 表示每个注意力头使用的隐藏向量的维度，并且满足：d=hidden_size=head_size * head_num。P 和 head num 应该需要满足整除关系，从另一方面来看，其实也保证 MHA 计算的数值正确性，每个 Head 本身的 hd 维度需要完整，才能保证 Attention score 是数值正确的。因此，Ulysses 图中的 d/P，实际上只是按照 head num 进行切分，并不会对每个 head 自身的 hd 再次切分。

输入 O 的切分

首先，对输入矩阵 O 进行切分，每个卡上获得[N/P,d]大小分块的数据，即 Local O。DeepSpeed-Ulysses 不对权重进行切分，每个卡上持有完整的 Wq/Wk/Wv 权重。因此，每个卡上的[N/P,d]分块，可以直接和当前卡上的 Wq/Wk/Wv 相乘，得到[N/P,d]大小的 Q,K,V 矩阵作为输出，即 Local Q/K/V。注意，这些 Q,K,V 矩阵的性质是：所有 Head 的部分 Q, K, V，即 N/P。

切分 O - Local O - Local Q/K/V

对 d 维度执行 All-To-All

经过上一个步骤，我们得到的 Local Q/K/V，每个卡只有持有其中 N/P 部分，但是 Attention 部分要求对完整的 N 进行注意力机制的计算，并且，最好同时能实现 Attention 的分布式运算。为此，我们需要使得每卡恰好持有一个 Attention Head 的 N 维度上完整的 Q/K/V 值。这要怎么实现呢？DeepSpeed 团队发现，此时，只要对 d 维度执行一个 All-To-All 操作，恰好能实现这个功能。

对 d 维度执行 All-To-All

All-To-All 这个通信方式同 All-Reduce 一样，是分布式训练中的 Collective functions，All-to-All 在每个进程向每个其他进程发消息的一部分，最后处理器拥有各个进程消息的一部分，这个操作可以理解成一个分布式矩阵转置操作。以下图为例，执行 All-To-All 后，rank 0 持有 4 个[N/P, d/P]的数据，分别属于 N 的不同部分，因此可以进一步合并成[N, d/P=head_size]。这正好是一个 head 所需要的在 N 维度上完整的 Q/K/V。

reshape/重排

Full Attention 计算

由于经过 All-To-All 后，每个 rank 都恰好拿到了一个 head 在 N 维度上完整的 Q/K/V，即 Qh, Kh 和 Vh，因此可以按照正常的 Attention 流程进行注意力机制计算，每个 rank 得到当前 Head 完整的 Oh。

Full Attention On each Rank

对 N 维度执行 All-To-All

每个 rank 得到当前 Head 完整的 Oh 后，我们发现，为了能够使得后续的和 Wo 相乘以及 MLP 模块能够直接执行序列并行，我们还需要将 Oh([N, d/P])变换成 Local O([N/P, d])的形式，而这恰好又可以用另一个 All-To-All 来完成，即，对 N 维度执行 All-To-All。得到的结果恰好就是 Local O 的形式，那么，可以直接把结果传给 Wo 以及 MLP，继续执行序列并行。

对 N 维度执行 All-To-All Part-1

我们对这个图，继续展开，可以看到原来蓝色系列的小方块[N/P, d/P]是分散到不同的 rank 上，经过 All-To-All 集合通信后，相同色系的分块都被聚合到同一个 rank 上。将这些[N/P, d/P]按照 d 维度 concat 在一起，就得到了 Local O 需要的在 d 维度上完整的数据分块[N/P, d]。

对 N 维度执行 All-To-All Part-2

乘以 Wo

我们知道，标准的 Attention 在通过 softmax(QK^T)*V 后，还需要将结果和 Wo 相乘，对 O 做一次线性变换。在对 N 维度执行 All-To-All 后，我们得到了 Local O，维度为[N/P, d]，在 d 维度上是完整的。而 Wo 矩阵在，在每个 rank 上都完整地保存了一份，维度为[d, d]。因此，我们可以直接用每个 rank 上的 local O 和 Wo 相乘，得到[N/P, d]。

乘以 Wo

MLP 计算

MLP 的计算很好理解，因为在得到 Local O [N/P, d]后，MLP 的计算不像 Attention 那样，需要涉及完整的维度 N。每个 Token 的计算是相互独立的。计算公式如下：

得到的 Local Z 依然是[N/P, d]的形式，因此可以直接与下一个 Transformer Block 的 Wq/Wk/Wv 相乘，得到 Local Q/K/V，从而进入下一次序列并行，直至算到 Loss 或者输出生成的 Token。

MLP 计算

至此，DeepSpeed-Ulysses Sequence Parallelism 的算法流程就已经讲完了。这里贴一下论文中 DeepSpeed-Ulysses 和同期 Megatron-LM SP+TP 对比的训练性能数据。不过后来 Megatron-LM 支持了 CP（Context Parallelism），因此，这个对比图也已经过时了。

DeepSpeed-Ulysses vs Megatron-LM SP+TP

0x04 通信量分析

本小节，对 DeepSpeed-Ulysses 和 Megatron-LM SP 的通信量进行简单的分析。本小节的内容，主要为归纳整理已有的一些分析结论，这些结论主要来源于 DeepSpeed-Ulysses[2]和 Megatron-LM SP[1]各自的论文，以及大佬们的一些文章[3] [4] [5]。本文主要关注 fwd 阶段。首先，贴一下 DeepSpeed-Ulysses 论文中给的结论。

序列并行通信复杂度

DeepSpeed-Ulysses 通信复杂度

从通信复杂度的角度看，DeepSpeed-Ulysses 的通信复杂度为 O(M/P)，Megatron-SP 的通信复杂度为 O(M)。其中，M 表示 Message，即每个 rank 需要发送的信息总量（只看 send）。对于所有 GPU 之间的 All-to-All 通信，当消息总大小为 M 时，每个链接传输的通信量为 M/P。对于具有隐藏层大小 d、序列长度 N 和并行度为 P（P 个 GPU）的 Transformer 模型，DeepSpeed-Ulysses 在 Attention 计算之前对 QKV project 进行 All-to-All 通信，其消息总大小为 3Nd（Q,K,V 各算一份）；在 Attention 得到每个 rank 的 Oh 后再对 N 维度进行一次 All-to-All 通信，消息大小为 Nd（这时只有 O 需要通信）。因此，DeepSpeed-Ulysses 序列并行，所需要的每个链接的集合通信量为 4Nd/P（或复杂度为 O(N/P)，4d 为常量）。值得注意的是，4Nd/P 和 P 有关系，因此当 N 和 P 按比例增加时，每个卡的通信量可以保持不变。也就是，对于 DeepSpeed-Ulysses 序列并行来说，在序列长度 N 非常大时，增加卡数可以降低通信量。这个结论非常重要。[2]

Megatron-LM SP 通信复杂度

相比之下，现有的方法如 Megatron-LM，无论 P 如何变化，其通信量都会随着信息量 M 线性增加，导致通信复杂度为 O(M)。例如，Megatron-LM 在每个 Transformer 层执行 2 次 All-Gather（每次通信量为 Nd），以及 2 次 Reduce-Scatter（每次通信量为 Nd）。因此，总通信量为 4Nd，但是，当 P 远大于 1 时，每次 All-Gather 和 Reduce-Scatter 通信量均为 Nd，这个值保持不变，不会随着 P 的增大而减少。因此，Megatron-LM SP 的每个链接通信量为 4Nd，这比 DeepSpeed-Ulysses 序列并行性的通信量大 P 倍。这意味着如果模型序列长度变长（N 变大），Megatron-LM SP 无法通过扩张卡的数量来减少单卡通讯量，那么单卡花在通讯上的时间就可能更多，进而降低了训练速度。[4]

eepSpeed-Ulysses 的缺点

不过虽然看起来，DeepSpeed-Ulysses 可以通过增加卡数来降低通信量，这个结论很 fancy，但实际上 P 是不能无限增加的，他会被 head_num 限制住。因为，在实际使用上，DeepSpeed-Ulysses 需要满足 head_num 能够被 P（卡数）整除的要求。由于 head_num 不可能无限增加，因此 P<=head_num。比如遇到 GQA 或者 MQA 情况，K、V 的 head_num 很小，此时，将导致 GPU 数目 P 也不能变得很大。并且，由于 DeepSpeed-Ulysses 使用了 All-to-All 通信，因此，对底层 NCCL 实现和网络硬件拓扑要求比较高。一般 All2All 跨机器涉及拥塞控制，带宽比较差。[5]

0x05 Ulysses + Zero3

序列并行本身的问题，是需要复制模型参数。DeepSpeed-Ulysses 的做法就是，采用 DS 家原厂出产的 Zero3，来降低冗余的模型参数存储。也就是，形成 Ulysses 序列并行+Zero3 模型并行的分布式方案。为什么是 Zero3，而不是 Zero 1 或 2。这是因为根据 Zero 论文中的定义，Zero-1 是对优化器做切分，Zero-2 是对优化器和梯度进行切分，Zero-3 则是对优化器、梯度和模型参数进行切分。因此，为了解决复制模型参数的问题，需要用 Zero-3 才行。Zero-3 模型并行，具体的操作为，将模型的权重分摊到各个卡上，比如现在有 4 卡，那么模型每一层的权重会被均匀地切分成 4 分，[L0,L1,L2,L3]，在 Ulysses fwd 的时候，在执行每一层实际的运算之前，会通过 All-Gather 来获取当前层完整的权重，也就是“逐层”的 All-Gather，L=Gather([L0,L1,L2,L3])，然后才进入实际的运算，算完后丢弃该层权重，开始获取下一层。当然，权重的获取和计算是可以通过流水线重叠的，这个细节就不再展开了。因此，Zero3 看着是模型并行（确实切分了模型）实际上是数据并行，因为在 fwd 的时候需要获取完整的权重再运算，运算时的权重没有被切分。

DeepSpeed-Ulysses + Zero3

0x06 总结

本文简要地讲解了 Megatron-LM 中常用的 Tensor Parallelism 和 Sequence Parallelism 的原理，并分析了使用 Sequence Parallelism 前后 per layer 激活值的显存占用，并注意到 Megatron-LM 的 Sequence Parallelism 需要结合它的 Tensor Parallelism 一起使用；本文还通过图解的方式，由浅入深地讲解了 DeepSpeed-Ulysses Sequence Parallelism 算法原理，根据论文的说明，DeepSpeed-Ulysses 巧妙地通过两个 All-To-All 实现了序列并行，并且对比 Megatron-LM 的 Sequence Parallelism，具有更低的通信量。不过在实际使用上，DeepSpeed-Ulysses 需要满足 head_num 能够被 P（卡数）整除的要求。由于 head_num 不可能无限增加，因此 P 实际上会被 head_num 限制住。

本人更多的技术笔记以及 CUDA 学习笔记，欢迎来 CUDA-Learn-Notes（CUDA Learn Notes with PyTorch）查阅。CUDA-Learn-Notes 包括了本人的 LLM/VLM 文章整理，以及对于 SGEMM/HGEMM/GEMV 等常见 CUDA Kernel 的示例实现，目前已经累计 2.5k+ stars，传送门：https://github.com/xlite-dev/...

CUDA Learn Notes with PyTorch
老样子，错误先更后改，持续更新...

参考

abReducing Activation Recomputation in Large Transformer Models
https://arxiv.org/pdf/2205.05198
abDeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models https://arxiv.org/pdf/2309.14509
图解大模型训练系列：序列并行 2，DeepSpeed Ulysses
https://zhuanlan.zhihu.com/p/...
ab 图解大模型训练系列：序列并行 1，Megatron SP
https://zhuanlan.zhihu.com/p/...
ab 大模型训练之序列并行双雄：DeepSpeed Ulysses & Ring-Attention
https://zhuanlan.zhihu.com/p/...

END

作者： DefTruth
来源：GiantPandaLLM

推荐阅读

欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式 AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。