[Arxiv 2019] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Introduction

• 作者设计了 Transformer 专用的模型并行策略 Megatron，利用张量并行 (Tensor Parallelism, TP) 有效降低了单卡显存需求和时延开销，但相比 DP/PP 仍然具有较大的通信开销，以 LLM 推理为例，一次推理时一个 Transformer block 内就需要做两次 all-reduce，在 decoding 阶段，通常需要使用 continuous batching 来提升 GPU 利用率，随着 batch size 增加，MHA 和 FFN 的 kernel 计算时延不会明显增加，而 all-reduce 的通信量却线性增加，相应的通信时延基本也线性增加，以至于 all-reduce 通信可能进一步成为瓶颈。此外，continuous batching 的方式会希望组合尽可能大的 batch size，也就意味着 GPU 上同一时间可能只有一个 CUDA 计算流，当 all-reduce 通信的时候并没有其他流的计算可以 overlap，导致通信开销难以掩盖。而 GPU 间的通信时延与 GPU 之间的互联方式有关，比如节点内采用 PCIE 或 NVLink 互联，跨节点采用 IB 网卡互联等会导致时延有数倍的差距，因此 TP 一般只会在单机内部使用

Model Parallel Transformers

Transformer layer

MLP

• (1) $Y=\text{GeLU}(XA)$ (GEMM + GeLU). 如果 $X$ 按列切分 $X=[X_1,X_2]$，$A$ 按行切分 $A=\left[\begin{array}{r}{A}_1\\ A_2\end{array}\right]$，那么 GEMM 可以很好地并行 $XA=X_1{A}_1+X_2{A}_2$，但由于 GeLU 为非线性函数，$\text{GeLU}(X_1{A}_1+X_2{A}_2)\ne \text{GeLU}(X_1{A}_1)+\text{GeLU}(X_2{A}_2)$，因此 GeLU 前需要插入 synchronization point. 而如果我们把 $A$ 按列切分 $A=[A_1,A_2]$ (Column Parallel) 就不需要在 GeLU 前就同步了：
  $$[Y_1,Y_2]=[\text{GeLU}(X{A}_1),\text{GeLU}(X{A}_2)]$$
• (2) $Z=\text{Dropout}(YB)$. 第二个线性层权重应该按行切分 $B=\left[\begin{array}{r}{B}_1\\ B_2\end{array}\right]$ (Row Parallel)，两个设备分别计算 $Z_1=Y_1{B}_1$ 和 $Z_2=Y_2{B}_2$，这样两个线性层的计算都分布在了不同设备上。然后通过 all reduce 即可得到计算结果 $Z_1+Z_2$，最后经过 Dropout 得到最终的输出
• Communication Overhead. 前向传播时 $g$ all-reduce $Z$ ($X$ 是每个 GPU 上都有的，因此无需 broadcast)；反向传播时 $f$ all-reduce $\nabla X$

Self-Attention

• out projection $Z=\text{Dropout}(YB)$ 的部分与 MLP 相同，下面只分析 $Y=\text{Self-Attention}(X)$. 对于多头自注意力层，每个头都可以放到一个单独的设备上并行计算 (当然，一个设备也可以负责多个头的计算，只要头数能被设备数整除即可保证负载均衡). $Q,K,V$ 均为 $k\times k$ 的矩阵，假设有 $h$ 个头，则 $Q,K,V$ 均可按列切为 $h$ 个 $k\times (k/h)$ 子矩阵，对应 $h$ 个自注意力头，分别放在不同设备上并行计算，每个自注意力头都输出 $N\times (k/h)$ 的子矩阵 $Y_i$
• Communication Overhead. 前向传播时 $g$ all-reduce $Z$；反向传播时 $f$ all-reduce $\nabla X$

Communication Overhead

• This enables us to perform all GEMMs in a simple transformer layer using only two all-reduces in the forward path and two in the backward path.

Input/Output embedding

• Input/Output embedding 共享权重，都是 $k\times V$ 的大矩阵 $E$，它们应该用相同的切法。作者将 $E$ 按列切 $E=[E_1,E_2]$，每个设备上仅有部分 token 的 embed.
• 对于 input embedding，由输入的 token 序列得到完整的 embed 序列需要进行 all-reduce ($g$ operator) (e.g. token 0 的 embed 存在设备 0 上，因此其余设备上 token 0 的 local embed 均为零向量，经过 all-reduce 后所有设备都能得到 token 0 的 embed).
• 对于 output embedding，parallel GEMM $[Y_1,Y_2]=[X{E}_1,X{E}_2]$ 加上 all-gather $Y=\text{all-gather}([Y_1,Y_2])$ 即可得到 logits. 但这样的话 all-gather 的通信量将是 $b\times s\times V$ ($b,s$ 分别为 batch size 和 seq len)，由于 $V$ 很大，因此通信量也很大，为了减少训练时的通信开销，可以先不对 logits 做 all-gather，而是对后续的 CE loss 做 all-gather. 具体来说，需要先由 logits 通过 softmax 得到 prob，也就是每个设备先对 $Y_i$ 计算 $\exp (Y_i)$，然后计算每行的和得到 $s_i\in {\mathbb{R}}^k$，对所有 $s_i$ 做 all-reduce 得到归一化分母 $s$，$\exp (Y_i)/s$ 即为 prob，再根据 GT label 即可计算出各个部分的 CE loss，再进行一次 all-reduce 即可得到损失

Other Layers

• 对于 dropout, layer normalization, residual connections 等其他部分，为了降低通信开销，Megatron 选择在每个设备上重复计算 (we maintain duplicate copies of layer normalization parameters on each GPU)

Summary

Layers	Model Parallel Method
Input Embedding	Row Parallel
Self-Attention	Column Parallel + Row Parallel
MLP	Column Parallel + Row Parallel
Output Embedding	Column Parallel

Sequence Parallelism (SP)

• 在原始的 Megatron 中，LN 和 dropout 需要在每个设备上重复计算，训练时它们的激活值也会占用相当大的内存，甚至超过权重和优化器状态，为了减少这部分内存占用，Megatron 团队在这篇论文里就提出了 Sequence Parallelism 对此做出了改进
• 具体来说，LN 和 dropout 都可以在序列维度上拆分到不同设备上去计算 (对于长序列训练比较友好)，相应地也需要在原有模型中插入通信原语；下图中 $g$ 和 $\bar{g}$ 是共轭的，$g$ 在前向传播时进行 all-gather，反向传播时进行 reduce-scatter，$\bar{g}$ 反之。总的通信开销为 4 次 all-gather 和 4 次 reduce-scatter，而 all-reduce 的通信开销为 all-gather/reduce-scatter 的两倍，因此通信开销不变！不仅节省了内存，还降低了单卡计算量，属于是一举两得了
  

---

• 为了提升通信效率，Megatron-LM 还做了一些通信优化：(1) Bulk Overlap. 在 SP (Sequence Parallel) 的反向传播阶段，可以将 all-gather (i.e., $\bar{g}$) 和计算输入的梯度进行隐藏 (i.e. overlap)，然后将 reduce-scatter (i.e., $g$) 和计算权重的梯度进行隐藏，而对于前向传播，由于通信和计算间存在依赖关系，因此无法 overlap；(2) Tensor Partitioning. 对于无法 overlap 的部分，可以使用 Tensor Partitioning 将一个大的矩阵乘法和集合通信操作，拆分成一系列小的矩阵乘法和集合通信操作，然后对更加细粒度的计算和通信进行流水线式的隐藏。当然，将一个 tensor 切分得太小，反而会影响实际性能，一般来说切成 4 份是比较常见的配置；(3) P2P Partitioning. 除了直接切分张量以外，我们还可以将集合通信操作拆分成一系列的 p2p 通信操作，例如 all-gather 操作可以拆分成 ring-based send/recv 通信，其中拆分后的通信和计算同样可以进行隐藏

Data Parallel (DP)

• 目前 Megatron 支持 ZeRO-1，即在 DDP 中实现 reduce-scatter 反向传递得到梯度，在 distributed optimizer 中实现 all-gather 优化器更新后的模型参数

Selective Activation Recomputation

• 在 LLM 训练中，中间激活值占用的显存为 $34bsh+5b{s}^{2}a$ (不使用 flash attention 的情况下)，其中 $5b{s}^{2}a$ 是为了保存 self-attention 中 softmax 的输入 $2b{s}^{2}a$ 以及 softmax 后 dropout 的输出 $2b{s}^{2}a$ 及其 mask $1b{s}^{2}a$，这部分激活值的显存开销实际上超过了全部激活值开销的一半 $5b{s}^{2}a>34bsh$，而重算这部分激活值所需的 FLOPs 又很小，因此 Selective Activation Recomputation 就是对这部分激活值做 recomputation，在仅引入极少额外 FLOPs 的情况下大幅降低激活值的内存占用；不过目前这个 trick 已被 flash attention 取代

References

• Shoeybi, Mohammad, et al. “Megatron-lm: Training multi-billion parameter language models using model parallelism.” arXiv preprint arXiv:1909.08053 (2019).
• Narayanan, Deepak, et al. “Efficient large-scale language model training on gpu clusters using megatron-lm.” Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. 2021.
• Korthikanti, Vijay Anand, et al. “Reducing activation recomputation in large transformer models.” Proceedings of Machine Learning and Systems 5 (2023): 341-353.
• code: https://github.com/NVIDIA/Megatron-LM
• 揭秘 LLM 推理：全面解析 LLM 推理性能的关键因素
• 详解 MegatronLM 序列模型并行训练 (Sequence Parallel)
• 大模型训练：Megatron-Core 中的通信优化