【学习笔记】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=\text{GeLU}(X_{(b\ell ),k}{A}_{k,k^{\prime }})\in {\mathbb{R}}^{(b\ell ),k^{\prime }}$$

对于矩阵 A 的分块方式

• 行分快
  • $X=\left[\begin{array}{c}{X}_1,X_2\end{array}\right],A=\left[\begin{array}{c}{A}_1\\ A_2\end{array}\right]$
  • $Y=\text{GeLU}(XA)=\text{GeLU}(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)$
    • $X_i{A}_i\in {\mathbb{R}}^{(b\ell ),k^{\prime }}$
• 列分快
  • $A=\left[\begin{array}{c}{A}_1,A_2\end{array}\right]$
  • $Y=\text{GeLU}(XA)=\text{GeLU}(X\left[\begin{array}{c}{A}_1,A_2\end{array}\right])=[\text{GeLU}(X{A}_1),\text{GeLU}(X{A}_2)]$
    • $X{A}_i\in {\mathbb{R}}^{b\ell ,k^{\prime }/2}$
  • 如果不同的 splits 放在不同的卡上，不同的卡需要维护全部的数据 $X$（数据未进行分块）

$$Z=\text{GeLU}(XA)B$$

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

• $B=\left[\begin{array}{c}{B}_1\\ B_2\end{array}\right]$

$$\begin{array}{rl}Z=& \text{GeLU}(XA)B\\ =& \left[\text{GeLU}(X{A}_1),\text{GeLU}(X{A}_2)\right]\left[\begin{array}{c}{B}_1\\ B_2\end{array}\right]\\ =& \text{GeLU}(X{A}_1)B_1+\text{GeLU}(X{A}_2)B_2\end{array}$$

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

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

mha

• 多头自注意力按照 num heads ($h$) 对 Q，K，V 三个 projection matrix 按列拆分 ($(k,k)\to (k,k/h)$ )
  • 对于 $O$：按行拆分
• 每个头的输出为 $Y_i=\text{softmax}\left(\frac{(X{Q}_i)(X{K}_i{)}^T}{\sqrt{d_k}}\right)V_i\in {\mathbb{R}}^{\ell ,k/h}$

$$[Y_1,Y_2]\left[\begin{array}{c}{B}_1\\ B_2\end{array}\right]=Y_1{B}_1+Y_2{B}_2$$

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\times v}=[E_1,E_2]$
  • 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]$
  • all-gather: (weight/tensor concat)
    • 存在通信的问题：$(b\times s)\times v$（$v$ 万级别的）
  • softmax：logits => probs
  • $X{E}_i\in {\mathbb{R}}^{(b\times s)\frac{v}{2}}$
  • $\text{rowsum}(\exp (X{E}_1))$, 长度为 $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 加一起；

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启动交互式配置。

---

[Pytorch 分布式] ring-allreduce 算法（scatter-reduce、allgather）以及 FSDP

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

code: https://github.com/chunhuizhang/pytorch_distribute_tutorials/blob/main/tutorials/3D-parallel/ring-allreduce.ipynb

之前探讨了DP、PP，这个要探讨SP的问题

Preliminary：FSDP（Fully Shared DP）

• all-gather/reduce-scatter

from IPython.display import Image

• N 张卡组成一个 ring 环，计算步数，2(N-1)
  • scatter-reduce: (N-1)，非标准 nccl
  • all-gather: (N-1)
• 3张卡，长度为6的向量加和为例；
  • input (each gpu model gradients):
    • [a0, a1 | a2, a3 | a4, a5] = [A0 | A1 | A2]
    • [b0, b1 | b2, b3 | b4, b5] = [B0 | B1 | B2]
    • [c0, c1 | c2, c3 | c4, c5] = [C0 | C1 | C2]
  • output (sync model gradients across gpus):
    • [a0+b0+c0, a1+b1+c1, a2+b2+c2, a3+b3+c3, a4+b4+c4, a5+b5+c5]
    • [A0 + B0 + C0 | A1 + B1 + C1 | A2 + B2 + C2]

这里要做的一个事情就是，三张卡上有三份不同的数据，现在要把它们加和起来并结算到某一张卡上去。

torch scatter reduce

• https://pytorch.org/docs/stable/generated/torch.Tensor.scatter_reduce_.html

import torch

src = torch.tensor([1., 2., 3., 4., 5., 6.])
index = torch.tensor([0, 1, 0, 1, 2, 1])
input = torch.tensor([1., 2., 3., 4.])
input.scatter_reduce(0, index, src, reduce="sum", include_self=True)

1+(1+3), 2+(2+4+6), 3+(5), 4
# tensor([5, 14, 8, 4])

src = torch.tensor([1., 2., 3., 4., 5., 6.])
index = torch.tensor([0, 1, 0, 1, 2, 1])
input = torch.tensor([1., 2., 3., 4.])
input.scatter_reduce(0, index, src, reduce="mean", include_self=True)

(1+(1+3))/3, (2+(2+4+6))/4, (3+(5))/2, 4
# tensor([1.667, 3.5, 4.0, 4])

上面的例子中，source是src = [1,2,3,4,5,6]，然后我们要把这些数据按照index中的索引进行scatter，scatter到input中去，然后做reduce操作（比如求和）

比如0这个位置要进来两个数据（1 和 3），也就是（1+1+3），其他位置是一样的，input里原先的数据也是要算进去的。

phase1： scatter reduce

减少通信量；先分块，以降低通信量，下面介绍的是

• [a0, a1 | a2, a3 | a4, a5] = [A0 | A1 | A2]
• [b0, b1 | b2, b3 | b4, b5] = [B0 | B1 | B2]
• [c0, c1 | c2, c3 | c4, c5] = [C0 | C1 | C2]
• scatter：data chunks，reduce：规约（降维）
  • nccl 是 reduce-scatter
  • 下面两步走是ring-allreduce的一算法
• step1
  • GPU0 =>(A2) GPU1 =>(B0) GPU2 =>(C1) GPU0
    • GPU0: A1 + C1, [A0, A1+C1, A2]
    • GPU1: B2 + A2, [B0, B1, B2+A2]
    • GPU2: C0 + B0, [C0+B0, C1, C2]
• step2
  • GPU0 =>(A1+C1) GPU1 =>(B2+A2) GPU2 =>(C0+B0) GPU0
    • GPU0: [C0+B0+A0, A1+C1, A2]
    • GPU1: [B0, A1+C1+B1, B2+A2]
    • GPU2: [C0+B0, C1, B2+A2+C2]

上面就是第一轮一个环状的传数据，第二轮也是环状传数据三张卡，第一轮只传了一个块，第二轮就两个块，两轮结束，需要计算的所有数据都有了，然后就是reduce/gather到一张卡上。

phase2: all-gather

gather再两步，就三张卡都有需要的数据了。

• S0: A0+B0+C0, S1: A1+B1+C1, S2: A2+B2+C2
• step1:
  • GPU0 =>(S0) GPU1 =>(S1) GPU2 =>(S2) GPU0
    • GPU0: [S0, …, S2]
    • GPU1: [S0, S1, …]
    • GPU2: […, S1, S2]
• step2:
  • GPU0 =>(S2) GPU1 =>(S0) GPU2 =>(S1) GPU0
    • GPU0: [S0, S1, S2]
    • GPU1: [S0, S1, S2]
    • GPU2: [S0, S1, S2]

下面是对上面两步走操作的图示例：

why ring-allreduce

• 高效的带宽利用率 (Efficient Bandwidth Utilization):
  • 分块传输: Ring-AllReduce 将需要同步的数据（例如梯度）分成多个小块（chunks）。
  • 流水线效应: 数据块在环上逐步传输和计算。一个 GPU 可以同时发送一个块给下一个节点，并从上一个节点接收另一个块。这种流水线方式使得 GPU 间的通信链路（如 NVLink 或网络带宽）能够持续被利用，而不是在等待整个大块数据传输完成。
  • 点对点通信: 每个 GPU 只需与其在环中的直接邻居通信。这使得算法可以充分利用现代 GPU 系统中高速的点对点连接（如 NVLink），避免了所有 GPU 都向一个中心点发送数据可能造成的拥塞。理论上，在 N 个 GPU 的环中，每个 GPU 在 Scatter-Reduce 和 All-Gather 阶段总共发送和接收的数据量大约是 2 * (N-1)/N * TotalDataSize，接近于最优值 2 * TotalDataSize。
• 均衡的通信负载 (Balanced Communication Load):
  • 在 Ring-AllReduce 中，每个 GPU 发送和接收的数据量大致相同，计算负载（Reduce 操作）也相对均衡地分布在各个步骤中。
  • 这避免了像基于树（Tree-based）的 All-Reduce 算法中可能出现的根节点通信瓶颈问题，因为在树形结构中，靠近根节点的 GPU 需要处理更多的数据聚合或分发任务。
• 避免中心瓶颈 (Avoids Central Bottleneck):
  • 与参数服务器（Parameter Server）架构或其他需要中心协调节点的同步方法不同，Ring-AllReduce 是完全去中心化的。没有单个节点会成为性能瓶颈或单点故障。
• 良好的可扩展性 (Good Scalability):
  • 虽然完成一次完整的 Ring-AllReduce 需要 2 * (N-1) 步（N 是 GPU 数量），延迟会随着 N 线性增加，但关键在于每个 GPU 的带宽需求基本保持不变（与 N 无关）。
    对于带宽是主要瓶颈的大规模系统（尤其是在传输大量梯度时），这种恒定的带宽需求使得 Ring-AllReduce 比那些带宽需求随节点数增加而增加的算法更具扩展性。

---

FSDP回顾

https://www.bilibili.com/BV1Kx4y187Te