【学习笔记】pytorch分布式

张量并行与 megtron-lm 及 accelerate 配置

https://www.bilibili.com/video/BV1TsWoe4E22

https://arxiv.org/abs/1909.08053

megtron-lm: 顾名思义针对 transformer 来做的优化
- 是 mp（论文题目），其实更多是tp（Tensor 张量内部做split）
- Transformer（intra layer parallel）
  - mlp
  - mha
  - embedding (input: wte, output: lm_head)
单卡做基线，没有通信的开销。存在划分，必然就存在通信。
集成进 accelerate
- accelerate 的几个 backends
  - deepspeed
  - fsdp
  - megtron-lm
- https://huggingface.co/docs/accelerate/usage_guides/megatron_lm

mlp

Y = GeLU ( X ( b ℓ ) , k A k , k ′ ) ∈ R ( b ℓ ) , k ′ Y=\text{GeLU}(X_{(b\ell),k}A_{k,k'})\in \mathbb R^{(b\ell),k'} Y=GeLU(X(bℓ),kAk,k′)∈R(bℓ),k′

对于矩阵 A 的分块方式

行分快
- X = [ X 1 , X 2 ] , A = [ A 1 A 2 ] X=\begin{bmatrix}X_1,X_2\end{bmatrix},A=\begin{bmatrix}A_1\\A_2\end{bmatrix} X=[X1,X2],A=[A1A2]
- Y = GeLU ( X A ) = GeLU ( X 1 A 1 + X 2 A 2 ) Y=\text{GeLU}(XA)=\text{GeLU}(X_1A_1+X_2A_2) Y=GeLU(XA)=GeLU(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)
  - X i A i ∈ R ( b ℓ ) , k ′ X_iA_i\in\mathbb R^{(b\ell),k'} XiAi∈R(bℓ),k′
列分快
- A = [ A 1 , A 2 ] A=\begin{bmatrix}A_1,A_2\end{bmatrix} A=[A1,A2]
- Y = GeLU ( X A ) = GeLU ( X [ A 1 , A 2 ] ) = [ GeLU ( X A 1 ) , GeLU ( X A 2 ) ] Y=\text{GeLU}(XA)=\text{GeLU}(X\begin{bmatrix}A_1,A_2\end{bmatrix})=[\text{GeLU}(XA_1),\text{GeLU}(XA_2)] Y=GeLU(XA)=GeLU(X[A1,A2])=[GeLU(XA1),GeLU(XA2)]
  - X A i ∈ R b ℓ , k ′ / 2 XA_i\in \mathbb R^{b\ell,k'/2} XAi∈Rbℓ,k′/2
- 如果不同的 splits 放在不同的卡上，不同的卡需要维护全部的数据 X X X（数据未进行分块）

Z = GeLU ( X A ) B Z=\text{GeLU}(XA)B Z=GeLU(XA)B

对于矩阵 B 自然进行行分块：

B = [ B 1 B 2 ] B=\begin{bmatrix}B_1\\B_2\end{bmatrix} B=[B1B2]

Z = GeLU ( X A ) B = [ GeLU ( X A 1 ) , GeLU ( X A 2 ) ] [ B 1 B 2 ] = GeLU ( X A 1 ) B 1 + GeLU ( X A 2 ) B 2 \begin{split} Z=&\text{GeLU}(XA)B\\ =&\left[\text{GeLU}(XA_1),\text{GeLU}(XA_2)\right]\begin{bmatrix}B_1\\B_2\end{bmatrix}\\ =&\text{GeLU}(XA_1)B_1 + \text{GeLU}(XA_2)B_2 \end{split} Z===GeLU(XA)B[GeLU(XA1),GeLU(XA2)][B1B2]GeLU(XA1)B1+GeLU(XA2)B2

最后对两张卡计算结果的加和是一种 all-reduce 的过程

关于all reduce可参考https://zhuanlan.zhihu.com/p/469942194，本质上是一个优化节点数据通信的算法，实现是比较容易的，阿里巴巴的ACCL

mha

多头自注意力按照 num heads ( h h h) 对 Q，K，V 三个 projection matrix 按列拆分 ( ( k , k ) → ( k , k / h ) (k,k)\rightarrow (k,k/h) (k,k)→(k,k/h) )
- 对于 O O O：按行拆分
每个头的输出为 Y i = softmax ( ( X Q i ) ( X K i ) T d k ) V i ∈ R ℓ , k / h Y_i=\text{softmax}\left(\frac{(XQ_i)(XK_i)^T}{\sqrt{d_k}}\right)V_i\in \mathbb R^{\ell,k/h} Yi=softmax(dk (XQi)(XKi)T)Vi∈Rℓ,k/h

[ Y 1 , Y 2 ] [ B 1 B 2 ] = Y 1 B 1 + Y 2 B 2 [Y_1,Y_2]\begin{bmatrix}B_1\\B_2\end{bmatrix}=Y_1B_1+Y_2B_2 [Y1,Y2][B1B2]=Y1B1+Y2B2

emb

如果词表数量是64000，嵌入式表示维度为5120，类型采用32 位精度浮点数，那么整层参数需要的显存大约为64000 × 5120 × 4 /1024/1024 = 1250MB，反向梯度同样需要1250MB，仅仅存储就需要将近2.5GB。
- [personal chatgpt] Llama2 7B vs. Llama3 8B （词表、attention 及 mlp）
wte: E H × v = [ E 1 , E 2 ] E_{H\times v}=[E_1,E_2] EH×v=[E1,E2]
- column-wise（v，vocab-size dimension）
- 1-50000: 1-25000, 25001-50000
- all-reduce (weight/tensor sum)
lm head: [ Y 1 , Y 2 ] = [ X E 1 , X E 2 ] [Y_1,Y_2]=[XE_1,XE_2] [Y1,Y2]=[XE1,XE2]
- all-gather: (weight/tensor concat)
  - 存在通信的问题： ( b × s ) × v (b\times s)\times v (b×s)×v（ v v v 万级别的）
- softmax：logits => probs
- X E i ∈ R ( b × s ) v 2 XE_i\in\mathbb R^{(b\times s)\frac v2} XEi∈R(b×s)2v
- rowsum ( exp ⁡ ( X E 1 ) ) \text{rowsum}(\exp(XE_1)) rowsum(exp(XE1)), 长度为 b s bs bs 的列向量，同理长度为 b s bs bs 的列向量，两个列向量 all-reduce 继续得到长度为 bs 的列向量

[0, 1, 25000, 25001]: input，不进行拆分
- 索引 E1 => 4*hidden_size，第3-4行为全0；
- 索引 E2 => 4*hidden_size，第1-2行为全0；
- 两个结果通过 all-reduce 加一起；

py 复制代码

import torch
import torch.nn.functional as F

torch.manual_seed(42)

A = torch.randn(5, 8)  # 5行12列的随机矩阵
"""
tensor([[ 1.9269,  1.4873,  0.9007, -2.1055,  0.6784, -1.2345, -0.0431, -1.6047],
        [-0.7521,  1.6487, -0.3925, -1.4036, -0.7279, -0.5594, -0.7688,  0.7624],
        [ 1.6423, -0.1596, -0.4974,  0.4396, -0.7581,  1.0783,  0.8008,  1.6806],
        [ 0.0349,  0.3211,  1.5736, -0.8455,  1.3123,  0.6872, -1.0892, -0.3553],
        [-1.4181,  0.8963,  0.0499,  2.2667,  1.1790, -0.4345, -1.3864, -1.2862]])
"""

A_1, A_2 = A.split(4, dim=1)

A_1
"""
tensor([[ 1.9269,  1.4873,  0.9007, -2.1055],
        [-0.7521,  1.6487, -0.3925, -1.4036],
        [ 1.6423, -0.1596, -0.4974,  0.4396],
        [ 0.0349,  0.3211,  1.5736, -0.8455],
        [-1.4181,  0.8963,  0.0499,  2.2667]])
"""

A_2
"""
tensor([[ 0.6784, -1.2345, -0.0431, -1.6047],
        [-0.7279, -0.5594, -0.7688,  0.7624],
        [-0.7581,  1.0783,  0.8008,  1.6806],
        [ 1.3123,  0.6872, -1.0892, -0.3553],
        [ 1.1790, -0.4345, -1.3864, -1.2862]])
"""

exp_A_1 = torch.exp(A_1)
exp_A_2 = torch.exp(A_2)

rowsum_exp_A_1 = torch.sum(exp_A_1, dim=1)
rowsum_exp_A_2 = torch.sum(exp_A_2, dim=1)

# all-reduce
rowsum = rowsum_exp_A_1 + rowsum_exp_A_2

rowsum.view(-1, 1)
"""
tensor([[17.2970],
        [10.2543],
        [19.1843],
        [14.4078],
        [17.8164]])
"""

exp_A_1 / rowsum.view(-1, 1)
"""
tensor([[0.3971, 0.2558, 0.1423, 0.0070],
        [0.0460, 0.5071, 0.0659, 0.0240],
        [0.2693, 0.0444, 0.0317, 0.0809],
        [0.0719, 0.0957, 0.3348, 0.0298],
        [0.0136, 0.1375, 0.0590, 0.5415]])
"""

exp_A_2 / rowsum.view(-1, 1)
"""
tensor([[0.1139, 0.0168, 0.0554, 0.0116],
        [0.0471, 0.0557, 0.0452, 0.2090],
        [0.0244, 0.1532, 0.1161, 0.2799],
        [0.2578, 0.1380, 0.0234, 0.0487],
        [0.1825, 0.0363, 0.0140, 0.0155]])
"""

torch.concat([exp_A_1 / rowsum.view(-1, 1), exp_A_2 / rowsum.view(-1, 1)], dim=1)
torch.allclose(softmax, torch.concat([exp_A_1 / rowsum.view(-1, 1), exp_A_2 / rowsum.view(-1, 1)], dim=1)) # True

accelerate megtron-lm config

https://huggingface.co/docs/accelerate/usage_guides/megatron_lm

Sequence Parallelism (SP): Reduces memory footprint without any additional communication.
- https://arxiv.org/pdf/2205.05198
  - （Megatron 3）
- Only applicable when using TP.
- It reduces activation memory required as it prevents the same copies to be on the tensor parallel ranks post all-reduce by replacing then with reduce-scatter and no-op operation would be replaced by all-gather.
- https://zhuanlan.zhihu.com/p/522198082
- LayerNorm和Dropout的计算被平摊到了各个设备上，减少了计算资源的浪费；
- LayerNorm和Dropout所产生的激活值也被平摊到了各个设备上，进一步降低了显存开销。

存在划分，必然就存在通信。在 Megatron1, 2 中，Transformer核的TP通信是由正向两个Allreduce以及后向两个Allreduce组成的。Megatron 3由于对sequence维度进行了划分，Allreduce在这里已经不合适了。为了收集在各个设备上的sequence parallel所产生的结果，需要插入Allgather算子；而为了使得TP所产生的结果可以传入sequence parallel层，需要插入reduce-scatter算子。在下图中，

所代表的就是前向Allgather，反向reduce scatter，

则是相反的操作。这么一来，我们可以清楚地看到，Megatron-3中，一共有4个Allgather和4个reduce-scatter算子。乍一看，通信的操作比Megatron-1 2都多得多，但其实不然。因为一般而言，一个Allreduce其实就相当于1个Reduce-scatter和1个Allgather，所以他们的总通信量是一样的。

如何配置？

在./.cache/huggingface/accelerate/default_config.yaml里修改。使用命令workspace accelerate launch启动交互式配置。