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

连理o

已于 2024-07-08 13:47:04 修改

阅读量446

点赞数

文章标签： Arxiv 2019

于 2023-09-27 14:35:36 首次发布

本文链接：https://blog.csdn.net/weixin_42437114/article/details/132867392

版权

7 篇文章 1 订阅

订阅专栏

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 一般只会在单机内部使用

在这里插入图片描述

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

在这里插入图片描述

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$

在这里插入图片描述

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 共享权重，都是 $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] = [XE_1, XE_2]$ 加上 all-gather $\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\R^k$ ，对所有 $s_i$ 做 all-reduce 得到归一化分母 $s$ ， $exp(Y_i)/s$ 即为 prob，再根据 GT label 即可计算出各个部分的 CE loss，再进行一次 all-reduce 即可得到损失

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

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

在原始的 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 的两倍，因此通信开销不变！不仅节省了内存，还降低了单卡计算量，属于是一举两得了

关注