分布式并行策略

总结

分布式并行策略总结

分布式策略总结

1. 数据并行

数据分多块、模型做复制放在设备上；需对梯度做AllReduce; 单卡容易装不下模型；
DP很慢、单机多卡，DDP多进程各进程单独优化、多机多卡；FSDP也能对梯度、优化器、模型状态做并行。
优点：简单易实现

2. 流水线并行

模型按层间划分多块，分发到不同设备上；缓解单卡压力，能训超大模型。
朴素流水线并行：每时刻只有1个GPU运行，GPU利用率极低 (bubble太大)。
Gpipe：同步
- minibatch切分为更多更小的microbatch，提升高并行度，重计算
- F-then-B模式，频繁流水线刷新，缓存太多micro中间变量和梯度，利用率低、内存要求高。
PipeDream：1F1B模式，解决内存高/刷新频繁等问题，异步。但存储多个版本模型，太多了。
PipeDream-2BW
- 最多只存储2个版本权重，每m个批次再生成一个版本权重，异步。极大降低内存使用。
PipeDream-FLUSH：同步，定期流水线刷新。
优点：适用于深层次模型。

3. 张量并行

模型按层内参数/Tensor划分，每个设备各自计算一部分，最终再汇总，有通信开销。 $Y = X A$
行并行：对模型参数按行划分，输入X需按列划分。列并行：对模型参数按列划分。以及MLP/Attention实践。
1D并行：按1个维度划分，显存占用大，通信成本高；
2D并行：把input和weight沿2个维度划分，降低显存占用；但通信高。
2.5D/3D并行：使用多个处理器降低通信， $p = q \cdot d \cdot d$ 。
优点：适用于训练非常大的模型。

3.5 3D并行 DP + PP + TP

3D并行是超大模型训练的必备。到千张卡时，3D并行比Zero3好。

4. 序列并行

输入长度为n，self-attention内存需求是 $o (n^{2})$ ；长序列数据增加activations等中间数据存储，显存高，限制训练能力
Clossal-AI：输入长度分割成多块；
Megatron-LM：把LayerNorm和Dropout的输入，按输出长度维度进行切分；其计算和激活值分配给不同设备，降低显存。

5. 混合并行

DP+PP
DP+PP+TP：3D并行

6. MoE并行

moe+数据并行：门和专家网络都会被复制，专家数量受单GPU大小限制。
moe+模型并行：门网络复制+专家网络放在不同设备，专家数量受限于GPU数量。

分布式并行策略选择

并行策略选择

单机单卡
- ZeRO + Offload CPU
- 启用ZeRO以内存为中心的平铺(MCT)：允许自动分割层顺序执行，减少了实时GPU运行参数数量，但不影响激活内存
单机多卡
- 模型能在单卡运行：DDP、ZeRO
- 不行不能在单卡运行：PP、ZeRO、TP
多机多卡
- ZeRO：无需修改模型
- PP+TP+DP：通信少，但需要修改模型
- 不推荐：PP+ZeRO-2/3，PP需要累计梯度，ZeRO会对梯度进行分块，这会影响效率。

背景

参考文章

常见的分布式并行策略

不并行示意图

数据 $x$ 、模型 $w$ 、结果 $o u t$ 不并行示意图：

梯度AllReduce

AllReduce

All Reduce = Reduce-Scatter + All-gather

All Reduce功能： 聚合梯度并广播回各节点，等同于以下2阶段。
Reduce-Scatter阶段
- 在每个GPU上，基于rank索引对rank相等的块 进行求和。
- GPU0: A0+B0+C0+D0，GPU1: A1+B1+C1+D1， GPU2: A2+B2+C2+D2
All-gather阶段
- 各GPU上求和后的梯度，再广播到其他相同位置的块，供所有GPU使用。
- GPU0: A,B,C,D, GPU01: A,B,C,D, GPU02: A,B,C,D, GPU3: A,B,C,D,
总结
- 各GPU按rank索引负责对相同块梯度求和，再广播回其他GPU。完成梯度聚合和广播

数据并行

基本概念

数据并行

背景：训练数据太大
核心方法
- 数据：🔑数据切分成多块，每块在一个设备上。
- 模型：🔑模型复制多份，各设备拥有完整且相同的模型副本。
- 梯度：各设备梯度不同，需对梯度进行累加合并，再广播到其他设备
  - 梯度AllReduce。
- 沿Batch维度做训练并行
缺点👿
- 每个设备存一个完整模型可能存不下，代价很大😞
- 梯度同步代价非常大

数据并行示意图

朴素数据并行(Pytorch DP)

1. 数据并行 (PyTorch DP，早期)

单进程多线程实现
- 梯度reduce到主卡，主卡更新参数 再广播到其他卡
缺点
- 💔很慢，只能单机多卡，不支持分布式，不支持APEX混合精度训练。
- ⚠️主卡(GPU0) 性能和通信开销是瓶颈，GPU利用率低且不均衡。
- 不支持模型并行

分布式数据并行(Pytorch DDP)

2. 分布式数据并行 (Pytorch DDP)

多进程实现
- 进程之间互相通信且只传梯度，大大节省通信开销。
每个进程独立优化
- 💡各自做 前向计算、反向传播、参数更新。
- ⭐各GPU之间 做梯度AllReduce、都得到平均梯度后，才开始反向。
DDP 相比 DP 更适合多机多卡，效率更高。
缺点：不支持模型并行👿。

完全分片数据并行(Pytorch FSDP)

3、完全分片数据并行 (Pytorch FSDP)

FSDP vs DDP

DP/DDP：只能做数据并行，不能做模型并行。
FSDP：数据并行训练方法，也能做模型并行(参数、梯度、优化器状态)，很像ZeRO3

FSDP

FSDP把状态分片到所有数据并行worker中，也可以把模型参数分片 卸载到CPU上。
模型层以嵌套方式用FSDP包装。
只有单个FSDP实例里的层做前后向计算才需把完整参数收集到设备
- 计算后将立即释放，做下一层计算。
- 不做计算时，可以把参数、梯度、优化器卸载到CPU。

优点

⭐节省峰值GPU内存，能做更大模型训练。

FSDP：

模型并行

📕核心思想

背景
- 每个GPU存放完整模型，存在冗余/开销大等多种问题。🔑
- 虽然FSDP可以缓解冗余，但对大模型来说，仅用数据并行难以进一步提升参数规模 🔑
模型并行
- 模型切分为多个部分，分布在多个设备上⭐。
两种模式
- 流水线并行：层间并行，对模型不同的Transformer层进行分割🌟
- 张量并行：层内并行，对Transformer层内进行分割⭐

模型-流水线并行/层间并行

层间并行：对不同的Transformer层进行分割并行

流水线并行思想和策略

🐱流水线并行思想

背景：模型太大 一个设备放不下
思想
- 💡模型按层分割成多块(阶段)，分发到不同设备上。
- 各设备之间以接力方式完成训练🔑。
  - 前向：每个设备把中间的激活 传给下一阶段
  - 反向：每个设备把tensor的梯度 回传上一阶段
缺点：👿容易出现设备空闲状态

🐸流水线并行策略

F-Then-B
- 先进行前向计算，再进行反向计算
- 缺点：缓存了多个micro-batch中间变量和梯度，GPU利用率不高
🔑1F1B
- ⭐前向和反向交叉并行。
- 及时释放不必要的中间变量👍，缓存activation数仅和stage数有关。
非交错式调度1F1B
- 第一阶段：热身，处理器进行不同数量的前向计算。
- 中间阶段：1F1B
- 最后阶段：完成后向计算。

Pipeline 并行示意图

4层网络切分到2个设备上
GPU0计算完T1和T2后，GPU0输出当做GPU1输入，GPU1 计算后面2层。

F-Then-B策略

所有F计算完成之后，才会计算B

1F1B策略

核心思想

F和B交替进行

示例

F42开始前：F41的反向B41已结束，可及时释放F41的变量和显存。
缓存activation数 仅和stage数有关

朴素流水线并行

1. 朴素流水线并行

核心思想

💥把模型按层间切分成多个Stage(部分)，每个Stage分配给一个GPU。
在每个Stage的边界处进行通信
并行策略：F-then-B

例子

前向：GPU1计算中间结果传给GPU2，GPU2计算模型最终结果
反向：GPU2计算loss，梯度输出传回GPU1，GPU1完成更新

\begin{array}{l} 网络目标: & output = L 4 (L 3 (L 2 (L 1 (i n p u t)))) \\ GPU1: & i n s = L 2 (L 1 (i n p u t)) \\ GPU2: & o u t p u t = L 4 (L 3 (i n s)) \end{array}

缺点

👿GPU利用率非常低
- 任意时刻只有1个GPU在工作，其余GPU在空闲(Bubble/空泡率)。
- 通信和计算没有交错：发送中间输出和梯度时，没有GPU执行操作。
设备之间复制数据的通信开销大
内存需求高
- 先执行前向的GPU会保存整个小批量缓冲的所有激活，直到结束
- 如果batch很大，内存会有问题

微批次-Gpipe流水线并行

2. 微批次流水线并行

🐰微批次核心

整体和朴素流水线一样，按层切分，但有区别
- 🔑把minibatch 划分为更多microbatch，一起训练 提高GPU利用率。
- 允许不同GPU 同时计算。
常见方法：GPipe、PipeDream。

2.1 Gpipe方法 (谷歌)

🐫核心

🔥提高模型并行度：mini-batch拆分为多个更小的micro-batch，送给GPU进行训练。
通过重计算降低单设备显存消耗峰值
并行策略：F-then-B

🔑重计算

⭐不保存中间层输出的激活值
⭐前向时记录每个算子的结果，反向时 重新计算出激活值 来计算梯度

⚠️缺点：

原因：minibatch切分成m份macrobatch后
🔑流水线刷新很频繁，降低硬件效率，导致空闲时间增加 ⏰。
- 流水线刷新/pipeline flush
  - 为确保权重一致 计算有效，定期 把当前状态进行同步更新。
需缓存m份activation，导致内存增加 👿
- 为什么需要缓存？
  - 很简单，前向的 中间结果activation 要被后向使用
- 重计算技术能解吗？
  - 即使重计算，前向激活也要等对应的后向计算完成后才能释放。

GPipe的Pipeline flush很频繁：

微批次-PipeDream流水线并行

2.2 PipeDream方法 (DeepSpeed)

🐴 核心

采用1F1B策略，解决缓存activation过多问题，缓存数量仅和stage有关(见1F1B图)。
目标：每个微批次 尽早完成后向计算--> 每个激活能早点释放--> 尽量减少激活保存时间

🛠️具体做法

1F1B策略，每台机器严格交替前向和后向计算。
- Megatron里对应参数virtual_pipeline_model_parrallel_size可以切分得更小。
保证每个GPU都有一个micro-batch数据在处理，流水线均衡，flush较少

💔缺点

使用权重存储方案，来保证相同输入的前后向权重版本一致
最多存d个权重版本，存储太多 💔
- 而Gpipe仅维护一个单一版本的模型权重。

PipeDream 变体

PipeDream-2BW

只维护2个版本模型权重 (double-buffered weights)，极大降低内存使用
每m个批次生成一个新的权重版本(m>=d，d为流水线深度)
缓冲：后向依赖于旧版本模型，新模型无法立即取代做缓冲用来未来使用。

PipeDream-Flush

同Gpipe，维护单权重、定期流水线刷新。
内存占用低于2BW，但吞吐量较低，以性能换取峰值内存。

Virtual Pipeline Parallel

Virtual Pipeline Parrallel

问题背景

虚拟流水线并行度，用于 Interleaved 调度，减少气泡。

基础PP：1F1B

GPU1：层 1- 4，GPU2：层5-8

Virtual PP

把原属于1张卡的层，切分成多块。比如8层，VPP=2，切分成16块。
- 卡1：块/[层 1-2]、[层 9-10]
- 卡2：[层 3-4]、[层 11-12]
卡1算完[1-2]，就可以扔给卡2，不用等[1-4]完全算完。
流水线流的很快、气泡更小。

各Pipeline 方法对比图

Gpipe 因为频繁Pipeline Flush导致空闲时间多

PipeDream 采用1F1B 空闲时间少

Gpipe和PipeDream-Flush的 Flush对比：

PipeDream-Flush占用内存最少，PipeDream-2BW 次之，GPipe会OOM。

PipeDream-2BW吞吐量最高，PipeDream-Flush次之，Gpipe吞吐量很差，Batch越大越体现。

流水线并行总结

1. 同步Pipeline并行🐱

朴素流水线并行 (非微批次, F-Then-B)
- 核心：每时刻只有1个GPU在运行，F-Then-B
- 缺点：GPU利用率极低💔
GPipe (微批次，F-Then-B)
- 核心：minibatch 切分为更小的macrobatch 给GPU提高并行度，使用重计算缓解内存。
- 缺点：频繁流水线刷新，导致硬件效率低和空闲时间高💔，缓存m份激活导致内存需求高
- 常见框架：Pytorch
PipeDream-Flush (1F1B)
- 核心：只存储单个权重，定期流水线刷新，以性能消耗降低峰值内存。
- 缺点：流水线刷新带来吞吐量变低。
- 常见框架：
  - DeepSpeed(非交错1F1B)、Megatron-LM(交错1F1B)、Colossal-AI(非交错+交错)

2. 异步Pipeline并行

PipeDream
- 核心：采用1F1B策略，解决GPipe内存高/刷新频繁等问题
- 缺点：采用模型权重存储方式，最多存储d个版本，存储太多导致内存高
PipeDream-2BW 🌟
- 核心：最多只存储2个版本权重，每m个批次再生成一个版本权重
- 优点：极大降低了内存使用

模型-张量并行/层内并行

基础思想

Tensor并行

思想
- 把Tensor沿特定维度 切分成n块💡，每个设备只持 $\frac{1}{n}$
- 每个设备各自计算一部分🔑，再通过通信串联起完整结果。
- 同时不影响图的计算正确性。
例子 C=AB计算
- $B = [B_{0} B_{1} \dots B_{n}]$
- 每个设备各自计算 $[A B_{0}, A B_{1}, \dots A B_{n}]$ ，最后再通信广播合并起来。
优点：节省了多设备间的梯度AllReduce 🔑
缺点：带来额外的通信开销⚠️ (每个设备做广播通信)

张量并行方式(行并行和列并行)

行并行和列并行

📕核心

X是输入，A是模型， $Y = X A$ ，Y是模型输出。
把模型切割成2部分，放在不同的GPU上进行运算。

🐱行并行

模型权重A 按行切分成2部分： $[\begin{matrix} A_{1} \\ A_{2} \end{matrix}]$
- 为保矩阵运算，需要把X按列分割成2部分，X的列数需要和A的行数相等。
计算示例
- GPU0： $X_{1} A_{1} = Y_{1}$ ， GPU1： $X_{2} A_{2} = Y_{2}$

X A = X \cdot \underset{按 行 切 分}{\underset{⏟}{[\begin{matrix} A_{1} \\ A_{2} \end{matrix}]}} = \underset{为 保 证 运 算}{\underset{⏟}{[X_{1} X_{2}]}} \cdot [\begin{matrix} A_{1} \\ A_{2} \end{matrix}] = X_{1} A_{1} + X_{2} A_{2} = Y_{1} + Y_{2} = Y

🐕列并行

模型权重A 按列切分 成2部分： $[\begin{matrix} A_{1} & A_{2} \end{matrix}]$
计算示例
- GPU0： $X A_{1} = Y_{1}$ ， GPU1： $X A_{2} = Y_{2}$

X A = [X] \underset{按 列 切 分}{\underset{⏟}{[\begin{matrix} A_{1} & A_{2} \end{matrix}]}} = [\begin{matrix} X A_{1} & X A_{2} \end{matrix}] = [\begin{matrix} Y_{1} & Y_{2} \end{matrix}] = Y

行并行示意图

列并行示意图

1D并行-Transformer示例

1D并行

核心

🔑把张量 按照某1个维度(1维矩阵)做切分，把参数 分到多个处理器上。
Megatron-LM：对权重划分后放到不同GPU上计算。

缺点

👿显存空间占用大
- 没有对activations进行划分，每个处理器需要存储整个Activations
😓通信成本高
- 每个处理器需与其他所有处理器做通信，成本随并行度增高而增高。

Transformer 1D并行细节

哪些部分需要并行

模型：输入embedding + N层Transformer + 输出层
Transformer：Self-Attention + FFN，核心是MHA和FFN做并行。

🐇FFN并行

结构：线性层A+ GeLU + 线性层B。
输入是X，过程：
- 先对A列切割：把X拷贝到2块GPU上，计算得到 $Y_{1}$ 和 $Y_{2}$ ，过激活函数后，作为B的输入。
- 再对B行切割：计算得到 $Z_{1}$ 和 $Z_{2}$ ，做1次AllReduce，相加结果得到Z⭐。
- B反向传播：把 $\frac{\partial L}{\partial Z}$ 传到2块GPU上，各自完成梯度计算
- A反向传播：各自求得 $\frac{\partial L}{\partial X_{1}}$ 、 $\frac{\partial L}{\partial X_{2}}$ 后，做1次AllReduce，相加得到 $\frac{\partial L}{\partial X}$ 。
$X A \to \underset{对A列切割}{\underset{⏟}{\begin{array}{lll} X A_{1} & = & Y_{1} \\ X A_{2} & = & Y_{2} \end{array}}} \to \underset{对B行切割}{\underset{⏟}{\begin{array}{lll} Y_{1} B_{1} & = & Z_{1} \\ Y_{2} B_{2} & = & Z_{2} \end{array}}} \to Z$

🐰MHA并行

结构：Q、K、V + 线性层。
输入是X，整体过程和MLP一样，过程是：
- 先对Q、K、V做列切割：每个头放到一块GPU上。
- 再对线性层做行切割：最终输出Z时，做1次Allreduce。
- 反向时，QKV梯度再做1次Allreduce算梯度。

MLP并行：

MHA并行：

1个Transformer前向和反向总共4个All-Reduce操作。

多维并行

由于1D并行activations存储多和多个处理器通信，带来的显存高和通信成本高等问题，1D并行无法满足大模型的训练。提出了2D、2.5D、3D等并行方案。

重要

1D并行

核心：按1D进行划分
缺点：
- 显存占用大：没有对activations进行划分，每个处理器需存储整个激活
- 通信成本高：不同处理器需要相互通信，成本随并行度增加而增加

2D并行

核心：把input和weight都沿2个维度进行划分 (SUMMA矩阵乘法算法)
优点：
- 降低内存占用：降低activations大小，可以提高batch size。
缺点：
- 可能带来更多通信开销

2.5D并行

核心
- 在2D-SUMMA基础上，使用P个处理器来减少通信
- $P = q \cdot q \cdot d$ 个处理器。
  - d=1时：2D并行
  - d=q时，3D并行
优点：降低了通信

3D 并行

序列并行

基本思路

序列并行

背景

传统self-attention的内存需求是输入长度 $n$ 的2次方 $o (n^{2})$
😓长序列数据会增加中间activations的内存使用量，限制了设备的训练能力。

ColosalAI-序列并行

核心：内存高效的并行方法。
- 输入序列分割成多个块，每块输入到相应设备中；并提出了环自注意力机制。
优点
- 打破单设备训练长度限制，不需要单个设备保存整个序列

MegatronLM-序列并行

核心
- 把LayerNorm和Dropout中的输入，按输入长度维度进行切分。
- 其产生的计算和激活值，都分配给不同的设备，降低了资源显存开销。
优点
- 消除部分冗余计算，用更少设备训练大模型

ColosalAI

Megatron-LM

混合并行技术

DP+PP

DP：Rank0、Rank1
- Rank1 看不见 GPU0，Rank0 也看不见 GPU-1。
PP：Rank0：GPU-0、GPU-2，Rank1：GPU-1、GPU-3
- GPU-0会把负载转移给GPU-2，GPU-1也会把负载转移给GPU-3
每个维度(DP组/PP组)至少需要2GPU，因此至少需4GPU

3D 并行 (DP+PP+TP)

DP+PP+TP

3D并行

数据+Pipeline+Tensor并行，至少需要8张卡。
- (注意：并非Tensor里的3D并行)

ZeRO-DP+PP+TP

ZeRO-DP + PP +TP

当ZeRO-DP与PP、TP结合时候，通常只启用ZeRO-1阶段。

ZeRO-1：优化状态。
ZeRO-2：优化状态、梯度。
- 若ZeRO-2和PP一起用
  - PP在每个micro-batch需在分片前聚合梯度，带来通信开销，影响PP效率。
  - PP需要累计梯度，ZeRO-2会对梯度进行分块，影响效率。
ZeRO-3：优化状态、梯度、模型参数。
- 本质是DP+MP的组合。

MoE 并行

MoE架构

MoE+数据并行

moe 数据并行

核心：数据并行，门网络和专家网络都会被复制放到设备上。
缺点：专家数量受单个计算单元/GPU大小限制。

MoE+模型并行

moe 模型并行

核心
- 门网络做复制，各个专家网络被独自放在各单元。
- 需引入额外的通信操作，保证更多的专家网络被同时训练。
缺点
- 专家网络数量受限于GPU数量。
- 额外的通信。

总结 ​

分布式并行策略总结 ​

分布式并行策略选择 ​

背景 ​

数据并行 ​

基本概念 ​

朴素数据并行(Pytorch DP) ​

分布式数据并行(Pytorch DDP) ​

完全分片数据并行(Pytorch FSDP) ​

模型并行 ​

模型-流水线并行/层间并行 ​

流水线并行思想和策略 ​

朴素流水线并行 ​

微批次-Gpipe流水线并行 ​

微批次-PipeDream流水线并行 ​

Virtual Pipeline Parallel ​

各Pipeline 方法对比图 ​

流水线并行总结 ​

模型-张量并行/层内并行 ​

基础思想 ​

张量并行方式(行并行和列并行) ​

1D并行-Transformer示例 ​

多维并行 ​

序列并行 ​

基本思路 ​

混合并行技术 ​

DP+PP ​

3D 并行 (DP+PP+TP) ​

ZeRO-DP+PP+TP ​

MoE 并行 ​

MoE+数据并行 ​

MoE+模型并行 ​

行业MoE并行方法 ​

总结

分布式并行策略总结

分布式并行策略选择

背景

数据并行

基本概念

朴素数据并行(Pytorch DP)

分布式数据并行(Pytorch DDP)

完全分片数据并行(Pytorch FSDP)

模型并行

模型-流水线并行/层间并行

流水线并行思想和策略

朴素流水线并行

微批次-Gpipe流水线并行

微批次-PipeDream流水线并行

Virtual Pipeline Parallel

各Pipeline 方法对比图

流水线并行总结

模型-张量并行/层内并行

基础思想

张量并行方式(行并行和列并行)

1D并行-Transformer示例

多维并行

序列并行

基本思路

混合并行技术

DP+PP

3D 并行 (DP+PP+TP)

ZeRO-DP+PP+TP

MoE 并行

MoE+数据并行

MoE+模型并行

行业MoE并行方法