分布式训练面试八股：复习笔记

这篇文章是我准备分布式训练方向面试时整理的复习笔记。核心思路是：每个框架都有它发家的点，学它那个代表性的点，其余细节可以暂时忽略。

本文覆盖的主题包括：

PyTorch DDP
Megatron：TP、SP、CP
DeepSpeed：ZeRO-1/2/3、Ulysses
MoE：DeepSpeed、DeepEP

Open Table of contents

PP（Pipeline Parallelism）
DP、DDP、ZeRO
TP（Tensor Parallelism）
SP、CP（Sequence Parallelism / Context Parallelism）
EP（Expert Parallelism）
小结
参考资料

PP（Pipeline Parallelism）

流水线并行的核心问题是 bubble。GPipe 的关键贡献有两个：

micro-batch：把一个 mini-batch 拆成更细的 micro-batch 往流水线里灌，降低 bubble 比例。
re-materialization（activation checkpointing）：前向时不保存全部激活值，反向时重新计算，用时间换显存。

一些常见概念：

mini-batch vs. micro-batch：mini-batch 是训练迭代一次用的总样本；micro-batch 是流水线上每个阶段同时处理的一小片。
VPP（Virtual Pipeline Parallelism / interleaving PP）：在多个模型副本之间交错 stage，进一步压缩 bubble。

面试观点

PP 其实「没啥好做的」，很多场景默认 PP = 1。理解 bubble 和 micro-batch 的 trade-off 即可。

DP、DDP、ZeRO

DP（Data Parallelism）

经典数据并行：每个 worker 保存一份完整模型，处理不同数据，梯度更新通过参数服务器或 All-reduce 聚合。适合单机多卡场景。

DDP（Distributed Data Parallel）

支持多机多卡。核心通信算子是 Ring All-reduce：

Ring All-reduce = Reduce-scatter + All-gather

两个阶段各走一圈 ring，总通信量约为 $2 \times$ 模型参数。

ZeRO（Zero Redundancy Optimizer）

ZeRO 的思想是用时间换空间、用通信换显存。形式上像模型并行，实质上仍是数据并行：需要完整数据时通信，算完不是自己的就丢掉。

训练时显存里的主要组成部分：

模型权重
优化器权重
激活值
KV Cache
CUDA Graph 等额外开销

其中必须存的：

模型权重（fp32）
优化器权重（fp32，一阶动量 + 二阶动量）

$3 \times 4 = 12$ 倍模型参数量；中间激活值一般用 fp16。

策略	切分对象	额外通信	显存节省
ZeRO-1	优化器状态	$2 \times$ 参数	优化器状态降至 $1/ N_{d}$
ZeRO-2	优化器状态 + 梯度	$2 \times$ 参数	优化器状态 + 梯度降至 $1/ N_{d}$
ZeRO-3	优化器状态 + 梯度 + 模型权重	$3 \times$ 参数	三者均降至 $1/ N_{d}$

哪些是「白送」的？

ZeRO-1 和 ZeRO-2 是开 DP 后几乎白送的，肯定要用。ZeRO-3 把模型权重也切了，通信从 $2 \times$ 涨到 $3 \times$ ，并非白送，需要谨慎评估。

ZeRO 还有一些变体：ZeRO-R、ZeRO-Offload、ZeRO-Infinity，面试时知道它们把状态进一步卸到 CPU/NVMe 即可。

TP（Tensor Parallelism）

张量并行（TP）纵向切分，在 DeepSpeed 里也叫模型并行（MP），即 MP = TP。

矩阵乘法：

X \in R^{b \times s \times h}, W \in R^{h \times h^{'}}, Y = X W \in R^{b \times s \times h^{'}}

切分 $W$ 有两种方式：列切分与行切分。

列切分

$W$ 按列切成 $W_{i} \in R^{h \times h^{'} / N}$ ，则：

Y_{i} = X W_{i} \in R^{b \times s \times h^{'} / N}

天然可算，最后把 $Y_{i}$ concat 起来。

行切分

$W$ 按行切成 $W_{i} \in R^{h / N \times h^{'}}$ ，则 $X$ 也需要对应切成 $X_{i} \in R^{b \times s \times h / N}$ ：

Y_{i} = X_{i} W_{i} \in R^{b \times s \times h^{'}}

最后把 $Y_{i}$ elementwise 相加。

MLP 层的 TP

Y = (GELU (X) \cdot A) \cdot B

其中 $X \in R^{b \times s \times h}$ ， $A \in R^{h \times h^{'}}$ ， $B \in R^{h^{'} \times h}$ 。

$A$ 只能按列切分：因为 GELU 不满足 elementwise 相加，不能行切分。
$B$ 相应地按行切分。

这样 $A$ 的输出正好对应 $B$ 的输入，中间省去一次通信。

f 和 g 在不同阶段的角色：

前向 f：copy 完整 $X$
前向 g：All-reduce 得到 $Z$
反向 g：copy 链式法则第一项（loss 对 $Z$ 的梯度）
反向 f：All-reduce 更新梯度

MLP 层前向 + 反向共 2 次 All-reduce，通信量为 $4 \times$ 模型参数。

Attention 层的 TP

Self-Attention 天生适配列切分：Q、K、V 按列切分，FFN 行切分。Attention + MLP 前向各一次 All-reduce，反向也各一次，共 4 次 All-reduce，通信量为 $8 \times$ 模型参数。

Embedding 层的 TP

Embedding 分为两部分：

Word Embedding $(V, h)$ ： $V$ 很大，按行切分。
Positional Embedding $(L, h)$ ：一般不切或按行切分。

切分后的每个 rank 保存一个 token 区间，根据 token_id 映射到对应 rank 取 embedding，最后 All-reduce 聚合完整结果。

输出层 Embedding 的 TP

输出层按列切分，算出的 logits All-reduce 聚合。

为什么要和输入层共用一套 Embedding？

为了保证输入层和输出层共享同一套词向量，输入层按行切分、输出层按列切分的设计需要配合好。

Cross-entropy 层的 TP

紧跟输出层 Embedding 的结果：

每部分先按行求和，再 All-reduce 得到完整求和，与 ground truth 算交叉熵。
每部分每行的 loss 再 All-reduce 得到总 loss。

这两次 All-reduce 都很轻：第一次是 batch size 量级，第二次只有 $N$ 量级。复杂度从 $b \cdot s \cdot V$ 降到 $b \cdot s + N$ 。

TP + DP 联动

反向传播时，DP 组的下一层不依赖上一层的梯度聚合结果。因此：

TP 组内做 All-reduce。
DP 组可以往下继续算。

二者有一定 overlap，可以写算子优化。

SP、CP（Sequence Parallelism / Context Parallelism）

Megatron SP

Megatron SP 把 TP 的输入和输出也做切分。此时 f 和 g 要改为 All-gather 或 Reduce-scatter。

Megatron 的 TP + SP 联动时，MLP 层只需要处理自己那部分，直接省掉了 TP 时的 All-reduce！

通信量没变，但白送了序列并行

输入时 All-gather、输出时 Reduce-scatter，等价于一次 Ring All-reduce 的通信量，但额外获得了序列并行。

Attention + FFN 的前向 + 反向，就是 4 次 All-gather + 4 次 Reduce-scatter，共 8 倍通信量。

DeepSpeed Ulysses

Ulysses 是序列并行：每个 rank 持有部分序列 + 对应 QKV。进入 Attention 计算时，变为所有序列 + 某个 head 的 QKV；输出再换回来。

在 Attention 计算前后各引入一次 All-to-all：

前向：QKV chunk 各做一次 All-to-all，Attention 输出再做一次 All-to-all。
反向：All-reduce 算梯度可以被镜像的第一次 All-to-all 掩盖。

Ulysses 通信量低：前向 4 次 All-to-all，反向 4 次 All-to-all，每次通信量为 $(N \cdot d) / P$ ，总通信量为：

\frac{8 N d}{P}

相比 Megatron 多了一个 $P$ 的除数，所以不单独随 $N$ 增加而增大。 $P$ 与注意力头数有关，取值有限，可理解为除以一个常数，但已经是很好的优化。

Ulysses + ZeRO-3 联动

在 Ulysses 之前先用 ZeRO-3 All-gather 得到完整权重，其余逻辑与单独 Ulysses 基本一致。

EP（Expert Parallelism）

GShard 最早把 MoE 应用在 Transformer 上，奠定了 EP 的基础。

MoE 层的基本计算

输入数据 $(S, M)$ ，其中 $S = b s$ ， $M = d$ 。Gate 为 $(M, E)$ ， $E$ 是专家数。经过 Gate 得到概率 $(S, E)$ ，每个 token 去往每个 expert 的概率。GShard 取 top-2 专家。

负载均衡策略

Expert buffer：为每个专家设定 capacity。
随机路由：top-2 专家以一定概率随机乱走。
Auxiliary Loss：辅助损失惩罚负载不均。

边界情况：

zero padding：expert buffer 没满。
drop tokens：要去的 expert buffer 满了，token 被丢弃。

实现里常见的 dispatch_mask、combine_weight，后面有机会再展开。

EP + DP 联动

非 MoE 部分：TP + DP。
MoE 部分：EP + DP。

先决定 EP，DP 就确定了。

假设 EP = 4，DP = 4，共 16 个 GPU。前向时：

All-to-all dispatch：把 token 发到对应 expert。
Expert 计算。
All-to-all combine：把结果发回原 rank。

反向时，ep_dp_group（维护同一专家的 rank）做 All-reduce 算梯度，更新专家权重。

EP + TP 联动

在 EP 基础上再做 TP。All-to-all 时 TP 可能会发重复数据，改进方法是：

把 $(E, C, M)$ 的 $C$ 维按 TP group 内的 rank 数切分，每个 rank 只带自己的部分去参加 EP group 的 All-to-all。这本质上是用时间换空间、用通信换显存——要知道所有 token 发哪去，在 TP group 内做 All-gather 即可。

MoE 层输入不能 drop tokens

此时 MoE 层输入是开了 TP 的，必须认真算完再 All-reduce 还原，所以 MoE 层接收输入时不能 drop tokens。

EP + TP + DP 联动

先考虑 EP + TP，剩下的就是 DP。

MoE 层输入由于前面开了 TP，TP 组内各 rank 输入一致。接下来：

EP 组内做一次 All-to-all，把 token 发到 experts。
Expert 计算。
TP 组 All-reduce 得到完整结果。
EP 组内再做一次 All-to-all，把结果发回，送到下一个非 MoE 层。

反向时，DP 组需要 All-reduce 算梯度，传回前面的非 MoE 层。

EP + TP + PP + DP 联动

一般 PP = 1（DeepSpeed 就是如此）。硬要考虑也可以，就是 EP + TP 的基础上再加 PP，但太抽象，个人认为没必要。

小结

PP：解决单卡放不下整个模型的问题，关注 bubble 与 micro-batch。
DP/DDP/ZeRO：数据并行的演进，ZeRO-1/2 几乎是白送，ZeRO-3 需要权衡通信。
TP：层内切分，All-reduce 是主要开销，适合单节点内高带宽互联。
SP/CP：序列并行，Ulysses 通过 All-to-all 把通信量除以 $P$ 。
EP/MoE：专家并行，解决 MoE 层负载均衡与大规模专家部署。

参考资料

以下文章是我整理本文时的主要参考，按主题分组：

流水线并行
- 知乎专栏 · 图解大模型训练之：流水线并行（Pipeline Parallelism），以 Gpipe 为例
数据并行与 ZeRO
张量并行
- 知乎专栏 · 图解大模型训练之：张量模型并行（TP），Megatron-LM
MoE 与专家并行
- 知乎专栏 · MoE/EP 并行
序列并行与上下文并行