大模型并行技术：从数据并行、张量并行到 ZeRO/FSDP

本文结合chatgpt生成

大模型训练和推理的核心矛盾很简单：模型越来越大，单张 GPU 的显存和算力远远不够。

以一个参数量为N的模型为例，如果参数使用 BF16/FP16，每个参数占 2 Bytes，那么单纯模型权重就需要：

\[M_{\text{param}} = 2N \text{ Bytes} \]

一个 70B 参数模型，仅权重就大约需要：

\[70 \times 10^9 \times 2 = 140 \text{ GB} \]

这还没有算训练时的梯度、优化器状态、激活值、通信 buffer 等。因此，无论是训练还是推理，大模型都必须依赖多 GPU、多机集群和并行技术。

大模型并行技术的本质可以概括成一句话：

把模型、数据、激活值、优化器状态或计算任务拆开，让多张 GPU 协同完成原本单卡无法完成的任务。

常见的大模型并行技术包括：

1. 数据并行，Data Parallelism，DP
2. 张量并行，Tensor Parallelism，TP
3. 流水线并行，Pipeline Parallelism，PP
4. ZeRO / FSDP 参数切分并行
5. 序列并行，Sequence Parallelism，SP
6. 专家并行，Expert Parallelism，EP

下面逐一展开。

1. 为什么大模型需要并行？

训练一个大模型时，GPU 显存主要被以下几部分占用：

\[M_{\text{total}} = M_{\text{param}} + M_{\text{grad}} + M_{\text{optimizer}} + M_{\text{activation}} + M_{\text{buffer}} \]

其中：

• \(M_{\text{param}}\)：模型参数
• \(M_{\text{grad}}\)：反向传播产生的梯度
• \(M_{\text{optimizer}}\)：优化器状态，例如 Adam 的一阶矩和二阶矩
• \(M_{\text{activation}}\)：前向传播保存下来的中间激活
• \(M_{\text{buffer}}\)：通信 buffer、临时计算 buffer 等

如果使用 Adam 优化器进行混合精度训练，一个参数可能对应：

• BF16/FP16 参数：2 Bytes
• BF16/FP16 梯度：2 Bytes
• FP32 master weight：4 Bytes
• Adam 一阶矩 (m)：4 Bytes
• Adam 二阶矩 (v)：4 Bytes

粗略估算，训练时每个参数可能需要：

\[2 + 2 + 4 + 4 + 4 = 16 \text{ Bytes} \]

所以一个 70B 模型训练时，光参数、梯度和优化器状态就可能需要：

\[70 \times 10^9 \times 16 = 1120 \text{ GB} \]

这还没算激活值。单张 80GB A100/H100 显然放不下。

因此，大模型训练的核心问题不是“要不要并行”，而是：

到底应该怎么切，切在哪里，通信代价有多大，显存能省多少，吞吐能提升多少？

2. 数据并行：Data Parallelism

数据并行是最经典、最容易理解的并行方式。

假设有 \(p\) 张 GPU，每张 GPU 都保存一份完整模型。训练时，把一个 global batch 切成 \(p\) 份，每张 GPU 处理其中一份。

如果 global batch size 为 \(B\)，GPU 数为 \(p\)，那么每张 GPU 处理的 local batch size 为：

\[B_{\text{local}} = \frac{B}{p} \]

每张 GPU 独立完成前向传播和反向传播，得到本地梯度：

\[g_i = \nabla_\theta L_i(\theta) \]

其中 \(i\) 表示第 \(i\) 张 GPU。

为了让所有 GPU 上的模型参数保持一致，需要对所有 GPU 的梯度做平均：

\[g = \frac{1}{p} \sum_{i=1}^{p} g_i \]

然后每张 GPU 使用相同的平均梯度更新参数：

\[\theta_{t+1} = \theta_t - \eta g \]

这个梯度同步过程通常通过 AllReduce （集合通信操作，把所有 GPU 上的数据做求和/平均，然后让每张 GPU 都拿到同一个结果）完成。

2.1 数据并行的优点

数据并行的优点是简单、通用、扩展性好。

如果单张 GPU 能放下完整模型，那么增加 GPU 数量可以提升训练吞吐。因为每张 GPU 处理不同数据，总体 batch size 变大。

2.2 数据并行的缺点

数据并行最大的问题是：每张 GPU 都必须保存完整模型、完整梯度和完整优化器状态。

也就是说，数据并行不能解决“大模型单卡放不下”的问题。

它解决的是：单卡能放下模型，但训练吞吐不够，希望用更多 GPU 处理更多数据。

所以纯数据并行适合中小模型，但对于 70B、175B 级别的大模型，仅靠数据并行不够。

2.3 数据并行的通信量

数据并行每一步都需要同步梯度。假设模型参数大小为 \(S\) Bytes，使用 ring AllReduce(每张 GPU 只和相邻 GPU 通信。数据被切成多个块，在环上传递并累加，最后每张 GPU 都得到完整结果。避免单个 GPU 成为瓶颈) 时，每张 GPU 的通信量大约为：

\[V_{\text{AllReduce}} = 2 \cdot \frac{p-1}{p} \cdot S \]

当 \(p\) 很大时，近似为：

\[V_{\text{AllReduce}} \approx 2S \]

也就是说，每一步每张 GPU 大概要通信两倍模型梯度大小的数据。

这也是为什么在大规模训练中，通信带宽、NCCL、网络拓扑非常重要。

---

3. 张量并行：Tensor Parallelism

数据并行是“每张 GPU 放完整模型，切数据”。

张量并行则是：

多张 GPU 共同保存一个模型，把模型内部的大矩阵乘法拆开计算。

Transformer 中最核心的计算是矩阵乘法，比如 Linear 层：

\[Y = XW \]

其中：

• \(X \in \mathbb{R}^{B \times d_{\text{in}}}\)
• \(W \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}\)
• \(Y \in \mathbb{R}^{B \times d_{\text{out}}}\)

当 \(W\) 非常大时，可以把 \(W\) 拆到多张 GPU 上。

张量并行常见有两种切法：

1. 列并行，Column Parallel
2. 行并行，Row Parallel

3.1 列并行 Linear

列并行是按输出维度切分权重矩阵。

原始矩阵：

\[W \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}} \]

切成 \(p\) 份：

\[W = [W_1, W_2, \dots, W_p] \]

其中：

\[W_i \in \mathbb{R}^{d_{\text{in}} \times \frac{d_{\text{out}}}{p}} \]

每张 GPU 计算：

\[Y_i = XW_i \]

最终输出：

\[Y = [Y_1, Y_2, \dots, Y_p] \]

列并行的特点是：每张 GPU 只计算一部分输出通道。

如果后续算子也能接受被切分的输出，那么这里可以暂时不通信。否则需要 AllGather 把完整 \(Y\) 拼回来。

---

3.2 行并行 Linear

行并行是按输入维度切分权重矩阵。

原始矩阵：

\[W \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}} \]

按行切分：

\[W = \begin{bmatrix} W_1 \\ W_2 \\ \vdots \\ W_p \end{bmatrix} \]

同时输入 \(X\) 也按隐藏维度切分：

\( X = [X_1, X_2, \dots, X_p] \)

每张 GPU 计算局部结果：

\[Z_i = X_i W_i \]

完整输出是所有局部结果求和：

\[Y = \sum_{i=1}^{p} Z_i \]

因此，行并行 Linear 通常需要一次 AllReduce。

---

3.3 Transformer MLP 如何做张量并行？

Transformer 里的 MLP 通常形如：

\[\text{MLP}(X) = \sigma(XW_1)W_2 \]

其中：

• \(W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ffn}}}\)
• \(W_2 \in \mathbb{R}^{d_{\text{ffn}} \times d_{\text{model}}}\)
• \(\sigma\) 是激活函数，比如 GELU、SiLU

由于 \(d_{\text{ffn}}\) 通常远大于 \(d_{\text{model}}\)，例如：

\[d_{\text{ffn}} \approx 4d_{\text{model}} \]

MLP 是 Transformer 中计算量很大的部分。

Megatron-LM 风格的张量并行通常这么切：

第一层 \(W_1\) 做列并行：

\[W_1 = [W_{1,1}, W_{1,2}, \dots, W_{1,p}] \]

每张 GPU 计算：

\[H_i = \sigma(XW_{1,i}) \]

第二层 \(W_2\) 做行并行：

\[W_2 = \begin{bmatrix} W_{2,1} \\ W_{2,2} \\ \vdots \\ W_{2,p} \end{bmatrix} \]

每张 GPU 计算：

\[Z_i = H_i W_{2,i} \]

最后做 AllReduce：

\[Y = \sum_{i=1}^{p} Z_i \]

这样设计的好处是：

1. \(W_1\) 输出本来就被切开，不需要立刻通信。
2. 激活函数 \(\sigma\) 是逐元素操作，可以在每张 GPU 的局部张量上独立完成。
3. \(W_2\) 之后只需要一次 AllReduce。

---

3.4 Attention 如何做张量并行？

Self-Attention 的核心计算为：

\[Q = XW_Q,\quad K = XW_K,\quad V = XW_V \]

然后：

\[\text{Attention}(Q,K,V) = \text{softmax} \left( \frac{QK^T}{\sqrt{d_k}} \right)V \]

多头注意力本身天然适合并行。假设总共有 \(h\) 个 attention heads，可以把不同 heads 分配到不同 GPU 上。

如果有 \(p\) 张 GPU，每张 GPU 负责：

\[\frac{h}{p} \]

个 attention heads。

每张 GPU 局部计算一部分 heads 的 \(Q,K,V\) 和 attention 输出，最后再通过输出投影层合并。

这种切法的优势是：

Attention head 之间在大部分计算阶段彼此独立，非常适合按 head 切分。

不过在输出投影 \(W_O\) 处，通常会涉及一次通信。

3.5 张量并行的优缺点

张量并行的优点：

1. 可以让多张 GPU 共同承载一个大层。
2. 适合 Transformer 中的大矩阵乘法。
3. 对推理和训练都很重要。
4. 可以降低单卡参数显存压力。

张量并行的缺点：

1. 通信频繁，通常每层都可能发生通信。
2. 对 GPU 间带宽要求高。
3. 更适合单机多卡或高速互联环境，例如 NVLink。
4. TP degree（Tensor Parallelism degree，张量并行度，一个层内的矩阵计算被切到多少张 GPU 上，TP degree 就是多少） 过大时，通信开销可能超过计算收益。

一个典型判断是：

张量并行适合“单层太大”或者“单卡放不下完整层”的情况，但代价是更高频的层内通信。

4. 流水线并行：Pipeline Parallelism

张量并行是切一个层里面的矩阵。

流水线并行则是切模型的层。

比如一个 Transformer 有 48 层，可以分到 4 张 GPU 上：

GPU0: Layer 0  - 11
GPU1: Layer 12 - 23
GPU2: Layer 24 - 35
GPU3: Layer 36 - 47

前向传播时，数据依次经过 GPU0、GPU1、GPU2、GPU3。

这就是流水线并行。

4.1 为什么需要流水线并行？

当模型层数很多时，即使单层能放下，完整模型也可能放不下。

流水线并行把不同层放到不同 GPU 上，使每张 GPU 只保存一部分层的参数。

如果总参数量为 \(N\)，流水线并行度为 \(p\)，理想情况下每张 GPU 保存：

\[\frac{N}{p} \]

的参数。

4.2 Pipeline Bubble

流水线并行的问题是会产生空泡，pipeline bubble。

假设有 4 个 pipeline stage，只有一个 batch 输入时，执行过程类似：

Time 1: GPU0 工作，GPU1 空闲，GPU2 空闲，GPU3 空闲
Time 2: GPU0 工作，GPU1 工作，GPU2 空闲，GPU3 空闲
Time 3: GPU0 工作，GPU1 工作，GPU2 工作，GPU3 空闲
Time 4: GPU0 工作，GPU1 工作，GPU2 工作，GPU3 工作
Time 5: GPU0 空闲，GPU1 工作，GPU2 工作，GPU3 工作
Time 6: GPU0 空闲，GPU1 空闲，GPU2 工作，GPU3 工作
Time 7: GPU0 空闲，GPU1 空闲，GPU2 空闲，GPU3 工作

前面填充流水线和后面排空流水线时，都有 GPU 空闲。

为了解决这个问题，会把一个 batch 切成多个 micro-batch。

假设 pipeline stage 数为 \(p\)，micro-batch 数为 \(m\)，粗略情况下，pipeline bubble 比例可以写成：

\[\text{Bubble Ratio} \approx \frac{p-1}{m+p-1} \]

当 \(m\) 越大时，bubble 越小。

例如 \(p=4\)，\(m=1\)：

\[\frac{4-1}{1+4-1} = \frac{3}{4} \]

bubble 很严重。

如果 \(p=4\)，\(m=16\)：

\[\frac{3}{16+3} = \frac{3}{19} \]

bubble 明显降低。

4.3 GPipe 与 1F1B

流水线并行常见调度方式有两类：

1. GPipe
2. 1F1B

GPipe 的方式是：

先把所有 micro-batch 的 forward 全部跑完，再统一跑 backward。

这会导致需要保存大量激活值，因为 backward 要等所有 forward 完成之后才开始。

1F1B 的意思是：

One Forward One Backward，即稳定阶段每做一个 forward，就做一个 backward。

1F1B 可以减少激活值驻留时间，从而降低显存压力。

4.4 流水线并行的优缺点

流水线并行的优点：

1. 可以按层切分超大模型。
2. 显著降低单卡参数显存。
3. 适合模型层数很多的场景。
4. 可以和数据并行、张量并行组合使用。

缺点：

1. 有 pipeline bubble。
2. 需要平衡每个 stage 的计算量。
3. 跨 stage 需要传输激活值和梯度。
4. 调度复杂，debug 成本高。
5. micro-batch 数增大可能影响优化器行为和显存。

流水线并行适合解决：

整个模型太深、太大，无法完整放在单卡上的问题。

---

5. ZeRO 与 FSDP：切参数、梯度和优化器状态

数据并行的问题是每张 GPU 都保存完整模型、完整梯度和完整优化器状态。

ZeRO 和 FSDP 的核心思想是：

既然数据并行中很多状态在每张 GPU 上都是重复的，那就把这些状态切开存。

ZeRO 全称是 Zero Redundancy Optimizer，核心目标是消除数据并行中的冗余显存。

5.1 ZeRO-1：切优化器状态

在普通数据并行中，每张 GPU 都保存完整优化器状态。

ZeRO-1 把优化器状态切分到不同 GPU 上。

如果优化器状态总大小为 \(M_{\text{optimizer}}\)，数据并行度为 \(p\)，那么每张 GPU 只保存：

\[\frac{M_{\text{optimizer}}}{p} \]

但是参数和梯度仍然是完整复制的。

ZeRO-1 的显存大致为：

\[M_{\text{ZeRO-1}} = M_{\text{param}} + M_{\text{grad}} + \frac{M_{\text{optimizer}}}{p} \]

5.2 ZeRO-2：切优化器状态和梯度

ZeRO-2 进一步把梯度也切分。

每张 GPU 保存：

\[\frac{M_{\text{grad}}}{p} \]

显存大致为：

\[M_{\text{ZeRO-2}} = M_{\text{param}} + \frac{M_{\text{grad}}}{p} + \frac{M_{\text{optimizer}}}{p} \]

相比 ZeRO-1，ZeRO-2 对显存节省更明显。

---

5.3 ZeRO-3：切优化器状态、梯度和参数

ZeRO-3 最激进，它连模型参数也切分。

每张 GPU 只保存完整参数的一部分：

\[\frac{M_{\text{param}}}{p} \]

因此显存大致为：

\[M_{\text{ZeRO-3}} = \frac{M_{\text{param}}}{p} + \frac{M_{\text{grad}}}{p} + \frac{M_{\text{optimizer}}}{p} \]

也可以写成：

\[M_{\text{ZeRO-3}} = \frac{ M_{\text{param}} + M_{\text{grad}} + M_{\text{optimizer}} }{p} \]

这就是 ZeRO-3 能训练超大模型的关键。

5.4 ZeRO-3 的通信过程

ZeRO-3 虽然省显存，但代价是通信更多。

因为每张 GPU 不再保存完整参数，所以在计算某一层之前，需要先通过 AllGather 收集该层参数。

前向传播某层时：

\[\text{AllGather}(\theta_i) \rightarrow \theta \]

拿到完整参数后计算：

\[Y = f(X; \theta) \]

计算完成后，可以释放不需要的完整参数，只保留本地 shard。

反向传播时，需要对梯度做 ReduceScatter：

\[\text{ReduceScatter}(g) \rightarrow g_i \]

5.5 FSDP 和 ZeRO 的关系

FSDP，全称 Fully Sharded Data Parallel，是 PyTorch 中常用的参数切分训练方案。

可以粗略理解为：

FSDP 是 PyTorch 生态里的 ZeRO-3 风格实现。

FSDP 会把参数、梯度、优化器状态切分到多个 GPU 上，并在需要计算某一层时动态 AllGather 参数。

FSDP 的关键思想是：

参数不必一直完整驻留在每张 GPU 上，只在计算需要时短暂 materialize(临时得到完整参数)。

---

5.6 ZeRO/FSDP 的优缺点

优点：

1. 极大降低显存占用。
2. 可以训练远超单卡显存的大模型。
3. 相比张量并行，编程模型更接近数据并行。
4. 适合大规模训练。

缺点：

1. 通信开销较大。
2. 频繁 AllGather / ReduceScatter。
3. 对通信带宽敏感。
4. 需要和 activation checkpointing、mixed precision、overlap communication 等技术配合。
5. 参数 shard 粒度、wrap 策略会影响性能。

ZeRO/FSDP 适合解决：

数据并行冗余太多，参数、梯度、优化器状态占用显存过大的问题。

---

6. 序列并行：Sequence Parallelism

序列并行是把序列长度维度切开。

假设 Transformer 的输入张量形状为：

\[X \in \mathbb{R}^{B \times S \times H} \]

其中：

• \(B\)：batch size
• \(S\)：sequence length
• \(H\)：hidden size

序列并行会沿着 \(S\) 维度切分：

\[X = [X_1, X_2, \dots, X_p] \]

每张 GPU 只处理一部分 token。

6.1 为什么需要序列并行？

随着长上下文模型的发展，sequence length 越来越长。

Attention 的复杂度通常是：

\[O(S^2 H) \]

激活显存也会随 \(S\) 增长显著增加。

如果上下文长度从 4K 增加到 32K，attention 矩阵规模会从：

\[4096^2 \]

变成：

\[32768^2 \]

增长倍数为：

\[\left(\frac{32768}{4096}\right)^2 = 64 \]

所以长上下文训练中，序列维度成为重要瓶颈。

6.2 序列并行的典型用途

序列并行经常和张量并行配合使用，尤其用于降低：

1. LayerNorm 激活显存
2. Dropout 激活显存
3. 残差连接相关激活显存
4. 长序列场景中的中间状态显存

它不是单独解决所有问题的银弹，而是大规模训练中的一块重要拼图。

---

7. 专家并行：Expert Parallelism

专家并行主要用于 MoE，Mixture of Experts，混合专家模型。

MoE 的核心思想是：

模型里有很多 expert，但每个 token 只激活其中少数几个 expert。

一个 MoE 层可以表示为：

\[y = \sum_{i \in \text{TopK}(r(x))} a_i E_i(x) \]

其中：

• \(E_i\)：第 \(i\) 个 expert
• \(r(x)\)：router，决定 token 分配给哪些 expert
• \(a_i\)：router 给出的权重
• \(\text{TopK}\)：选择得分最高的 K 个 expert

如果总共有 \(E\) 个 expert，每个 token 只激活 \(K\) 个 expert，那么理论上计算量只和 \(K\) 有关，而参数量可以随 \(E\) 增大。

7.1 MoE 为什么需要专家并行？

MoE 模型的 expert 数量可能很多。单张 GPU 放不下所有 expert，或者放得下但计算效率不高。

专家并行会把不同 expert 放到不同 GPU 上。

例如：

GPU0: Expert 0, 1
GPU1: Expert 2, 3
GPU2: Expert 4, 5
GPU3: Expert 6, 7

每个 token 经过 router 后，被发送到对应 expert 所在的 GPU。

这个 token dispatch 过程通常需要 All-to-All 通信。

7.2 专家并行的瓶颈

专家并行的瓶颈主要有两个：

第一，All-to-All 通信。

token 分发到不同 expert，本质上是多对多通信。相比 AllReduce，All-to-All 更复杂，对网络拓扑更敏感。

第二，负载不均衡。

如果大量 token 都被 router 分配给少数 expert，那么这些 expert 所在 GPU 会非常忙，而其他 GPU 可能空闲。

为了解决负载不均衡，通常会加入 load balancing loss，例如：

\[L_{\text{balance}} = E \sum_{i=1}^{E} f_i p_i \]

其中：

• \(E\)：expert 数量
• \(f_i\)：分配给第 \(i\) 个 expert 的 token 比例。实际被分配到专家 i 的 token 比例，属于“硬路由结果”。
• \(p_i\)：router 对第 \(i\) 个 expert 的平均概率。router 对专家 i 给出的平均路由概率，属于“软路由倾向”。

等号右边乘 E 的主要作用是：归一化损失尺度，使它不随专家数量变化。

如果不乘 \(E\)，均衡状态下：

\[\sum_{i=1}^{E}f_ip_i = E\cdot\frac{1}{E^2} =\frac{1}{E} \]

这个损失鼓励 token 更均匀地分配到不同 expert 上。

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

在真实的大模型训练中，往往不是只用一种并行，而是组合使用。

最经典的是 3D 并行：

\[\text{Total GPUs} = D \times T \times P \]

其中：

• \(D\)：Data Parallel size
• \(T\)：Tensor Parallel size
• \(P\)：Pipeline Parallel size

例如有 64 张 GPU，可以配置为：

\[D = 8,\quad T = 4,\quad P = 2 \]

则：

\[8 \times 4 \times 2 = 64 \]

含义是：

• 每 4 张 GPU 做一个 tensor parallel group
• 每 2 个 stage 做 pipeline parallel
• 一共有 8 份 data parallel replica

8.1 三种并行分别解决什么问题？

可以这样理解：

并行方式	切分对象	主要解决的问题
数据并行 DP	数据 batch	提升吞吐
张量并行 TP	层内矩阵	单层太大，矩阵乘法太大
流水线并行 PP	模型层	模型层数太多，整体放不下
ZeRO/FSDP	参数、梯度、优化器状态	数据并行冗余显存太大
序列并行 SP	sequence 维度	长上下文激活显存太大
专家并行 EP	MoE expert	expert 数量太多

一句话总结：

DP 切数据，TP 切矩阵，PP 切层，ZeRO/FSDP 切训练状态，SP 切序列，EP 切专家。

---

9. 推理阶段的并行有什么不同？

训练和推理都需要并行，但关注点不同。

训练关注：

1. 显存能不能放下参数、梯度、优化器状态、激活值
2. 反向传播通信开销
3. AllReduce / AllGather / ReduceScatter
4. 计算和通信能否 overlap
5. 集群吞吐

推理关注：

1. 权重能不能放下
2. KV Cache 能不能放下
3. prefill 吞吐
4. decode 延迟
5. batch 调度
6. 多请求并发
7. Tensor Parallel 对单 token decode 的通信影响

9.1 推理中的 KV Cache 显存

大模型推理时，每生成一个 token，都需要缓存历史 token 的 key/value。

KV Cache 大小大致为：

\[M_{\text{KV}} = 2 \times L \times B \times S \times H_{\text{kv}} \times \text{bytes} \]

其中：

• 2 表示 K 和 V
• \(L\)：层数
• \(B\)：batch size
• \(S\)：sequence length
• \(H_{\text{kv}}\)：KV hidden size
• \(\text{bytes}\)：每个元素占用字节数，例如 FP16 为 2 Bytes

可以看到，KV Cache 随 batch size 和 sequence length 线性增长。

这也是为什么推理系统中会出现 PagedAttention、continuous batching、KV cache manager 等技术。

9.2 推理中的 Tensor Parallel

推理部署 70B 模型时，单张 GPU 往往放不下完整权重，所以 Tensor Parallel 很常见。

例如 70B BF16 权重大约 140GB。如果有 4 张 80GB GPU，可以用 TP=4，把权重切到 4 张卡上。

理想情况下，每张 GPU 保存：

\[\frac{140}{4} = 35 \text{ GB} \]

但实际还要考虑 KV Cache、临时 buffer、框架开销等。

推理中 TP 的问题是 decode 阶段 batch 较小时，单步计算量不大，但每层仍可能需要通信。因此，TP degree 不是越大越好。

10. 常见通信原语

理解大模型并行，必须理解几个通信原语。

10.1 AllReduce

AllReduce 会把所有 GPU 上的数据求和或平均，然后把结果广播回所有 GPU。

数学上：

\[y_i = \sum_{j=1}^{p} x_j,\quad i=1,2,\dots,p \]

所有 GPU 最后都得到相同的 \(y_i\)。

典型用途：

• 数据并行梯度同步
• 行并行 Linear 输出求和

10.2 AllGather

AllGather 会收集所有 GPU 上的数据 shard，并让每张 GPU 都得到完整数据。

如果每张 GPU 有：

\[x_i \]

AllGather 后每张 GPU 得到：

\[[x_1, x_2, \dots, x_p] \]

典型用途：

• ZeRO-3 / FSDP 前向前收集参数
• 张量并行中恢复完整激活

10.3 ReduceScatter

ReduceScatter 可以理解为 AllReduce 后再切分。

先求和：

\[y = \sum_{i=1}^{p} x_i \]

再把 \(y\) 切成 \(p\) 份，每张 GPU 拿一份。

典型用途：

• ZeRO/FSDP 梯度切分
• 数据并行优化通信

10.4 All-to-All

All-to-All 是每张 GPU 都向其他所有 GPU 发送不同数据。

典型用途：

• MoE token dispatch
• expert parallel

All-to-All 通信复杂度较高，容易成为 MoE 训练和推理的瓶颈。

11. 如何选择并行策略？

选择并行策略时，可以按下面的问题逐步判断。

11.1 单卡能否放下完整模型？

如果单卡能放下完整模型，优先考虑数据并行。

如果单卡放不下完整模型，需要 TP、PP、ZeRO/FSDP。

11.2 是单层太大，还是层数太多？

如果单层矩阵太大，比如 hidden size 很大，适合张量并行。

如果模型层数很多，适合流水线并行。

11.3 是训练，还是推理？

训练时需要考虑：

\[M_{\text{param}} + M_{\text{grad}} + M_{\text{optimizer}} + M_{\text{activation}} \]

推理时主要考虑：

\[M_{\text{param}} + M_{\text{KV}} + M_{\text{buffer}} \]

所以训练更关注 ZeRO/FSDP、activation checkpointing、optimizer state sharding。

推理更关注 Tensor Parallel、KV Cache、batching、低延迟。

11.4 通信带宽是否足够？

张量并行通信频繁，最好放在高速互联的 GPU 内部，例如同一台机器的 NVLink 内。

流水线并行通信频率相对低一些，可以跨机器，但要注意 stage 间传输激活值。

数据并行通常跨机器更常见，因为它主要每 step 同步梯度。

一种常见经验是：

TP 放在机内，PP 可以跨机，DP 扩展到多机。

当然，真实系统还要结合网络拓扑、模型结构、batch size 和框架实现。

---

12. 一个例子：训练 70B 模型该怎么并行？

假设要训练一个 70B Transformer，有 64 张 GPU。

可以考虑：

\[D \times T \times P = 64 \]

一种可能配置：

\[D=8,\quad T=4,\quad P=2 \]

含义：

• TP=4：每个层内矩阵切到 4 张 GPU
• PP=2：模型层切成 2 个 pipeline stage
• DP=8：整个 TP+PP 组合复制 8 份处理不同数据

如果还使用 ZeRO-1 或 ZeRO-2，可以进一步减少优化器状态和梯度显存。

如果使用 FSDP，则可能采用另一种策略，让参数、梯度、优化器状态在数据并行组内切分。

实际配置没有唯一答案，需要综合考虑：

1. 模型参数量
2. hidden size
3. 层数
4. batch size
5. sequence length
6. GPU 显存
7. NVLink / InfiniBand 带宽
8. 框架支持程度
9. 训练稳定性
10. pipeline bubble

---

13. 大模型并行技术概要

如果有人问：“你了解大模型并行技术吗？”

一个比较好的回答结构是：

大模型并行主要是为了解决单卡显存和算力不足的问题。常见方法有数据并行、张量并行、流水线并行、ZeRO/FSDP、序列并行和专家并行。数据并行切 batch，通过 AllReduce 同步梯度；张量并行切层内矩阵，比如 Megatron 中 MLP 第一层列切、第二层行切；流水线并行按层切模型，但会有 pipeline bubble，需要 micro-batch 和 1F1B 调度；ZeRO/FSDP 切 optimizer state、gradient 和 parameter，降低数据并行冗余；推理侧还要重点考虑 KV Cache、Tensor Parallel 和 batching。选择并行策略时，需要根据模型是否单卡放得下、单层是否过大、层数是否过多、通信带宽是否足够来决定。

这个回答既覆盖全局，又能落到通信和显存分析上。

14. 总结

可以用一句话总结：

数据并行切数据，张量并行切矩阵，流水线并行切层，ZeRO/FSDP 切训练状态，序列并行切 token 维度，专家并行切 expert。

再进一步，需要理解每种切分方式带来的通信代价：

技术	主要通信
数据并行	AllReduce 梯度
张量并行	AllReduce / AllGather
流水线并行	stage 间激活和梯度传输
ZeRO/FSDP	AllGather 参数，ReduceScatter 梯度
序列并行	Gather / Scatter 序列维度激活
专家并行	All-to-All token 分发