Contents
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
=
GeLU
(
X
A
)
Y=\text{GeLU}(XA)
Y=GeLU(XA) (GEMM + GeLU). 如果
X
X
X 按列切分
X
=
[
X
1
,
X
2
]
X=[X_1,X_2]
X=[X1,X2],
A
A
A 按行切分
A
=
[
A
1
A
2
]
A=\left[\begin{aligned}A_1\\A_2\end{aligned}\right]
A=[A1A2],那么 GEMM 可以很好地并行
X
A
=
X
1
A
1
+
X
2
A
2
XA=X_1A_1+X_2A_2
XA=X1A1+X2A2,但由于 GeLU 为非线性函数,
GeLU
(
X
1
A
1
+
X
2
A
2
)
≠
GeLU
(
X
1
A
1
)
+
GeLU
(
X
2
A
2
)
\text{GeLU}(X_1A_1+ X_2A_2) \neq \text{GeLU}(X_1A_1)+\text{GeLU}(X_2A_2)
GeLU(X1A1+X2A2)=GeLU(X1A1)+GeLU(X2A2),因此 GeLU 前需要插入 synchronization point. 而如果我们把
A
A
A 按列切分
A
=
[
A
1
,
A
2
]
A=[A_1,A_2]
A=[A1,A2] (Column Parallel) 就不需要在 GeLU 前就同步了:
[ Y 1 , Y 2 ] = [ GeLU ( X A 1 ) , GeLU ( X A 2 ) ] [Y_1, Y_2] = [\text{GeLU}(XA_1), \text{GeLU}(XA_2)] [Y1,Y2]=[GeLU(XA1),GeLU(XA2)] - (2) Z = Dropout ( Y B ) Z=\text{Dropout}(YB) Z=Dropout(YB). 第二个线性层权重应该按行切分 B = [ B 1 B 2 ] B=\left[\begin{aligned}B_1\\B_2\end{aligned}\right] B=[B1B2] (Row Parallel),两个设备分别计算 Z 1 = Y 1 B 1 Z_1=Y_1B_1 Z1=Y1B1 和 Z 2 = Y 2 B 2 Z_2=Y_2B_2 Z2=Y2B2,这样两个线性层的计算都分布在了不同设备上。然后通过 all reduce 即可得到计算结果 Z 1 + Z 2 Z_1+Z_2 Z1+Z2,最后经过 Dropout 得到最终的输出
- Communication Overhead. 前向传播时 g g g all-reduce Z Z Z ( X X X 是每个 GPU 上都有的,因此无需 broadcast);反向传播时 f f f all-reduce ∇ X \nabla X ∇X
Self-Attention
- out projection Z = Dropout ( Y B ) Z=\text{Dropout}(YB) Z=Dropout(YB) 的部分与 MLP 相同,下面只分析 Y = Self-Attention ( X ) Y=\text{Self-Attention}(X) Y=Self-Attention(X). 对于多头自注意力层,每个头都可以放到一个单独的设备上并行计算 (当然,一个设备也可以负责多个头的计算,只要头数能被设备数整除即可保证负载均衡). Q , K , V Q,K,V Q,K,V 均为 k × k k\times k k×k 的矩阵,假设有 h h h 个头,则 Q , K , V Q,K,V Q,K,V 均可按列切为 h h h 个 k × ( k / h ) k\times (k/h) k×(k/h) 子矩阵,对应 h h h 个自注意力头,分别放在不同设备上并行计算,每个自注意力头都输出 N × ( k / h ) N\times(k/h) N×(k/h) 的子矩阵 Y i Y_i Yi
- Communication Overhead. 前向传播时 g g g all-reduce Z Z Z;反向传播时 f f f all-reduce ∇ X \nabla X ∇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 × V k\times V k×V 的大矩阵 E E E,它们应该用相同的切法。作者将 E E E 按列切 E = [ E 1 , E 2 ] E=[E_1,E_2] E=[E1,E2],每个设备上仅有部分 token 的 embed.
- 对于 input embedding,由输入的 token 序列得到完整的 embed 序列需要进行 all-reduce ( g g 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 ] [Y_1, Y_2] = [XE_1, XE_2] [Y1,Y2]=[XE1,XE2] 加上 all-gather Y = all-gather ( [ Y 1 , Y 2 ] ) Y = \text{all-gather}([Y_1, Y_2]) Y=all-gather([Y1,Y2]) 即可得到 logits. 但这样的话 all-gather 的通信量将是 b × s × V b\times s\times V b×s×V ( b , s b,s b,s 分别为 batch size 和 seq len),由于 V V V 很大,因此通信量也很大,为了减少训练时的通信开销,可以先不对 logits 做 all-gather,而是对后续的 CE loss 做 all-gather. 具体来说,需要先由 logits 通过 softmax 得到 prob,也就是每个设备先对 Y i Y_i Yi 计算 exp ( Y i ) \exp(Y_i) exp(Yi),然后计算每行的和得到 s i ∈ R k s_i\in\R^k si∈Rk,对所有 s i s_i si 做 all-reduce 得到归一化分母 s s s, exp ( Y i ) / s \exp(Y_i)/s exp(Yi)/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
- 在原始的 Megatron 中,LN 和 dropout 需要在每个设备上重复计算,训练时它们的激活值也会占用相当大的内存,为了减少这部分内存占用,Megatron 团队在这篇论文里就提出了 Sequence Parallelism 对此做出了改进
- 具体来说,LN 和 dropout 都可以在序列维度上拆分到不同设备上去计算,相应地也需要在原有模型中插入通信原语;下图中
g
g
g 和
g
ˉ
\bar g
gˉ 是共轭的,
g
g
g 在前向传播时进行 all-gather,反向传播时进行 reduce-scatter,
g
ˉ
\bar g
gˉ 反之。总的通信开销为 4 次 all-gather 和 4 次 reduce-scatter,而 all-reduce 的通信开销为 all-gather/reduce-scatter 的两倍,因此通信开销不变!不仅节省了内存,还降低了单卡计算量,属于是一举两得了
References
- Shoeybi, Mohammad, et al. “Megatron-lm: Training multi-billion parameter language models using model parallelism.” arXiv preprint arXiv:1909.08053 (2019).
- Korthikanti, Vijay Anand, et al. “Reducing activation recomputation in large transformer models.” Proceedings of Machine Learning and Systems 5 (2023).
- code: https://github.com/NVIDIA/Megatron-LM
- 揭秘 LLM 推理:全面解析 LLM 推理性能的关键因素