总结

分布式并行策略总结

分布式策略总结

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=XA$
• 行并行：对模型参数按行划分，输入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$、结果$out$ 不并行示意图：

数据并行

基本概念

数据并行

• 背景：训练数据太大

• 核心方法

  • 数据：🔑数据切分成多块，每块在一个设备上。
  • 模型：🔑各设备需拥有完整且相同的模型副本，模型复制多份。
  • 梯度：各设备梯度不同，需对梯度进行累加合并，再广播到其他设备
    • 梯度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把这些状态分片到所有数据并行worker中，还可以选择把模型参数分片卸载到CPU上。
• FSDP
  • 模型层以嵌套方式用FSDP包装。
  • 只有单个FSDP实例里的层做前后向计算才需把完整参数收集到设备，计算后将立即释放，做下一层计算。
  • 不做计算时，可以把参数、梯度、优化器卸载到CPU。
• 优点
  • ⭐节省峰值GPU内存，能做更大模型训练。

FSDP：

AllReduce ReduceScatter All-gather

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，完成梯度聚合和广播

模型并行

模型并行

📕核心思想

• 背景
  • 每个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}{ll}\text{网络目标:}& \text{output}=L4(L3(L2(L1(input))))\\ \text{GPU1:}& ins=L2(L1(input))\\ \text{GPU2:}& output=L4(L3(ins))\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进行训练。
• 通过重计算降低单设备显存消耗峰值Re-materialization
• 并行策略：F-then-B

🔑重计算

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

⚠️缺点：

• 原因：minibatch切分成m份macrobatch后

• 🔑流水线刷新很频繁pipeline flsuh，降低硬件效率，导致空闲时间增加 ⏰。

  • 流水线刷新：为确保权重一致性、计算有效性，定期把当前状态进行同步或更新。

• 需缓存m份activation，导致内存增加 👿

  • 为什么需要缓存？
    • 很简单，前向的中间结果activation要被后向使用
  • 重计算技术能解吗？
    • 即使重计算，前向activation也要等对应的后向计算完成后才能释放。

GPipe的Pipeline flush很频繁：

微批次-PipeDream流水线并行

2.2 PipeDream方法 (DeepSpeed)

🐴 核心

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

🛠️具体做法

• 1F1B策略，每台机器严格交替前向和后向计算
• 保证每个GPU都有一个micro-batch数据在处理，流水线均衡，flush较少

💔缺点

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

🏇变体

• PipeDream-2BW
  • 只维护2个版本模型权重 (double-buffered weights)，极大降低内存使用
  • 每m个批次生成一个新的权重版本(m>=d，d为流水线深度)
  • 缓冲：后向依赖于旧版本模型，新模型无法立即取代做缓冲用来未来使用。
• PipeDream-Flush
  • 同Gpipe，维护单权重、定期流水线刷新。
  • 内存占用低于2BW，但吞吐量较低，以性能换取峰值内存。

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

PipeDream 采用1F1B 空闲时间少

Gpipe和PipeDream-Flush的 Flush对比：

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

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

流水线并行对比总结

流水线并行总结

1. 同步Pipeline并行🐱

• 朴素流水线并行 (非微批次)
  • 核心：每时刻只有1个GPU在运行
  • 缺点：GPU利用率极低💔
• GPipe⭐
  • 核心：minibatch切分为更小的macrobatch给GPU，提高并行度；使用重计算缓解内存。
  • 缺点：频繁流水线刷新，导致硬件效率低和空闲时间高💔，缓存m份激活导致内存需求高
  • 常见框架：Pytorch
• PipeDream-Flush ⭐
  • 核心：只存储单个权重，定期流水线刷新，以性能消耗降低峰值内存。
  • 缺点：流水线刷新带来吞吐量变低。
  • 常见框架：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\cdots B_n]$
  • 每个设备各自计算$[A{B}_0,A{B}_1,\cdots A{B}_n]$，最后再通信广播合并起来。

• 优点：节省了多设备间的梯度AllReduce 🔑

• 缺点：带来额外的通信开销⚠️ (每个设备做广播通信)

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

行并行和列并行

📕核心

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

🐱行并行

• 模型权重A按行切分成2部分：$\left[\begin{array}{c}{A}_1\\ A_2\end{array}\right]$
  • 为保矩阵运算，需要把X按列分割成2部分，X是列数需要和A的行数相等。
• 计算示例
  • GPU0：$X_1{A}_1=Y_1$， GPU1：$X_2{A}_2=Y_2$

$$XA=X\cdot \underset{\mathrm{按}\mathrm{行}\mathrm{切}\mathrm{分}}{\underbrace{\left[\begin{array}{c}{A}_1\\ A_2\end{array}\right]}}=\underset{\mathrm{为}\mathrm{保}\mathrm{证}\mathrm{运}\mathrm{算}}{\underbrace{[X_1{X}_2]}}\cdot \left[\begin{array}{c}{A}_1\\ A_2\end{array}\right]=X_1{A}_1+X_2{A}_2=Y_1+Y_2=Y$$

🐕列并行

• 模型权重A按列切分成2部分：$\left[\begin{array}{cc}{A}_1& A_2\end{array}\right]$
• 计算示例
  • GPU0：$X{A}_1=Y_1$， GPU1：$X{A}_2=Y_2$

$$XA=[X]\underset{\mathrm{按}\mathrm{列}\mathrm{切}\mathrm{分}}{\underbrace{\left[\begin{array}{cc}{A}_1& A_2\end{array}\right]}}=\left[\begin{array}{cc}X{A}_1& X{A}_2\end{array}\right]=\left[\begin{array}{cc}{Y}_1& Y_2\end{array}\right]=Y$$

行并行示意图

列并行示意图

1D并行-Transformer示例

1D并行

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}$。
  $$XA\to \underset{\text{对A列切割}}{\underbrace{\begin{array}{lll}X{A}_1& =& Y_1\\ X{A}_2& =& Y_2\end{array}}}\to \underset{\text{对B行切割}}{\underbrace{\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 并行

Pytorch

序列并行

基本思路

提示

背景

• 传统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)

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

ZeRO-DP+PP+TP

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

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

MoE 并行

MoE架构

MoE+数据并行

moe 数据并行

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

MoE+模型并行

moe 模型并行

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

行业MoE并行方法