第 21 章 分布式训练与并行

第 19、20 章都在单个 GPU 内部做优化。当模型本身有几千亿参数、或单卡吞吐远不能满足需求时，唯一的出路是把工作切开分给多个设备——代价是设备之间必须通信，而通信比计算贵得多。

学习目标

读完本章，你应当能够：

1. 说出驱动并行化的四大约束（显存容量、显存带宽、算力、通信），并判断某场景真正卡在哪一项。
2. 写出六种集合通信原语的语义与通信量（每个元素被搬运的次数），并把 All-Reduce 拆解为 Reduce-Scatter + All-Gather。
3. 对比 NVLink/NVSwitch 与 PCIe、InfiniBand 与以太网的带宽数量级，解释为什么高频通信的并行必须留在 NVLink 域内。
4. 描述 DP/DDP 的梯度同步流程，说明通信量为何与参数量同阶，以及 DDP 的 bucket 与通信计算重叠机制。
5. 推导 Megatron 风格张量并行"第一个线性层按列切、第二个按行切"的方案，说明为什么这样只需一次 All-Reduce，并给出每层前向/反向的通信次数。
6. 推导流水线并行的 bubble 比例公式，估算 1F1B 与虚拟流水线的收益。
7. 量化估算 ZeRO 三个阶段的显存节省，说明它与 FSDP 的关系。
8. 解释专家并行（EP/MoE）的稀疏激活、All-to-All 通信、容量因子与负载均衡。
9. 用决策树为给定模型规模与集群拓扑选出并行策略。

前置知识检查

21.1 为什么需要并行：四大约束

单卡做不了大模型训练，原因不止一个。分清楚是哪一项在卡，决定了该用哪种并行：

关键认识：前三项是"收益"，第四项是"成本"。并行化的全部工程内容，就是在"获得前三项收益"与"付出第四项成本"之间找平衡点。

这也解释了为什么并行策略的选择不只取决于模型，还取决于集群拓扑：同样的 8 卡 TP，放在 NVLink 全连接域内和放在 8 台机器上，效果能差一个数量级。这一点会在 21.3 与 21.5 中反复出现。

21.2 集合通信原语

并行训练的一切通信都可以归结为几个集合通信原语（collective communication primitives）。理解它们的关键不是名字，而是每个元素被搬运了几次。

设 N 为设备数（rank 数），D 为单个分片的数据量（元素数或字节数）。以下分析都基于环（ring）算法的朴素实现。

21.2.1 六种原语的语义

21.2.2 All-Reduce = Reduce-Scatter + All-Gather

All-Reduce 可以精确拆成两步，这两步恰好是"归约并分片"与"汇聚并广播"：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

All-Reduce（目标：每张卡都得到所有卡数据之和）
        │
        ├── 第 1 步  Reduce-Scatter
        │     各卡的数据被归约，但每张卡只拿到其中一个分片的结果
        │     通信量 ≈ D
        │
        └── 第 2 步  All-Gather
              各卡把自己持有的那个分片广播出去，人人拿到完整结果
              通信量 ≈ D
        ────────────────────────────
        合计 ≈ 2D          ← 与直接 All-Reduce 相同

为什么这个拆解很重要：如果某一个后续步骤本来就不需要完整结果（比如 ZeRO-2 里每张卡只需要自己负责的那片梯度来更新自己负责的那片参数），那么第一步 Reduce-Scatter 之后就可以停下——省掉整整一半通信。ZeRO/FSDP 的通信效率正建立在这个拆解上。

21.2.3 Ring All-Reduce 的思路

把 N 张卡连成一个环，每张卡只与左右两个邻居通信：

1
2
3
4
5
6
7

Ring All-Reduce（以梯形拼接为例）
  阶段 A（Reduce-Scatter，共 N−1 轮）：
     第 k 轮，每张卡把自己的一个分片发给右邻居、同时从左邻居收一个分片并累加
     N−1 轮之后：每张卡恰好持有"某一个分片的全局和"
  阶段 B（All-Gather，共 N−1 轮）：
     第 k 轮，每张卡把已算好的分片沿环传播，同时收下上游的
     N−1 轮之后：每张卡持有全部 N 个分片的全局和

优点：每张卡的收发量都是 2D(N−1)/N，与 N 无关；带宽利用率高（每个方向都不空闲）。缺点：延迟为 2(N−1) 轮，N 很大时延迟会成为主要开销——这正是"跨机 All-Reduce 慢"的机制来源之一。工程上用 树形 / 分层（hierarchical）All-Reduce 来减轻：机内先做一轮，机间再做一轮，把跨机通信的轮数降下来。

21.3 通信硬件：为什么 TP 必须放在同机

数量级上的结论很清晰：NVLink 比 PCIe 快约 7 倍，比 InfiniBand 快约 9 倍，比百 G 以太网快约 36 倍。

因此一条硬性工程规则：

这也解释了为什么实际大规模训练的标准形态是：机内 TP（8 卡满 NVLink）+ 机间 PP 与 DP。

21.4 数据并行：DP 与 DDP

21.4.1 流程

数据并行的做法最直观：每张卡放一份完整模型，各喂不同的数据。

1
2
3
4
5

第 t 步：
  ① 每张卡用自己的 micro-batch 做前向 + 反向，得到本地梯度 g_r
  ② 对全部 g_r 做一次 All-Reduce（求平均）
  ③ 每张卡用同一个平均梯度更新自己那份参数
     ⟹ 各卡参数始终完全一致

注意第 ③ 步为什么必须是"同一个梯度"：如果各卡用各自不同的梯度更新，参数就会分叉，下一轮的前向就失去了可比性。保证参数一致性是 DP 的正确性基础，也是一切通信开销的来源。

21.4.2 通信量与参数量同阶

第 ② 步的通信量：梯度张量大小为 P（参数量）个元素，All-Reduce 通信量约 2P（元素数）。所以

DP 每步通信量 ≈ 2 × 参数量

这是 DP 的核心特征：通信量与参数量同阶，与 batch size 无关。 推论很重要——增大全局 batch size 是"免费"降低 DP 通信占比的手段：数据多了，计算量线性增加，但同步一次的成本不变。

工程上还有两个补充手段：

• 梯度累积（gradient accumulation）：把 K 个小 micro-batch 的梯度先本地累加，再每 K 步同步一次。等效于把全局 batch 放大 K 倍，通信频率降到 1/K。代价是优化器的有效更新次数减少，学习率需要相应调整。
• DDP 的 bucket 与重叠：反向传播是逐层产生梯度的。DDP 不等全部梯度算完，而是把梯度按参数分成若干 bucket（桶），某个桶一填满就立刻发起 All-Reduce，与后续层的反向计算并行进行。

1
2
3
4
5
6
7
8

反向过程（从后往前）：
  层 L   反向计算 → 梯度入 bucket 1
  层 L-1 反向计算 → 梯度入 bucket 1   ← bucket 1 满，立刻发 All-Reduce
  层 L-2 反向计算 ─────────────────┐
  层 L-3 反向计算 ─────────────────┤ 与 bucket 1 的通信重叠
  层 L-4 反向计算 → 梯度入 bucket 2 ┘
  ...
  ⟹ 通信时间被隐藏在计算时间之后，端到端只看到 max(计算, 通信)

bucket 大小的取舍：太小则每次通信的消息小、启动开销占比高、带宽打不满；太大则第一个桶要等很久才满，重叠窗口变小、尾部还会剩下一个无法隐藏的通信尾巴。实践中常取几十 MB 级别，并按"反向产生的顺序"排列，让靠后产生的梯度（反向最先算出的那些层）先进桶。

21.5 张量并行：Megatron 风格的层内切分

数据并行不减少单卡显存（每卡一份完整模型），所以模型大到单卡放不下时必须切层内。张量并行（tensor parallelism, TP）把权重矩阵本身切开。

21.5.1 MLP 的列切与行切

考虑一个两层 MLP（先忽略激活函数）：Y = GeLU(X·A)·B，其中 X ∈ R^{n×d}、A ∈ R^{d×h}、B ∈ R^{h×d}。

天真做法（按列切 A）：把 A 按列切成 N 份 A = [A₁, A₂, …, A_N]，每张卡算 Z_r = GeLU(X·A_r)，其中 Z_r ∈ R^{n×h/N}。这没问题，因为 GeLU 是逐元素的，GeLU([X A₁, X A₂]) = [GeLU(X A₁), GeLU(X A₂)]——分块计算后再拼接，与整体计算完全等价。

但接下来要用 B。把 B 按行切成 B = [B₁; B₂; …; B_N]（与 Z 的列切分对应），则

Y = GeLU(X·A)·B = [Z₁, Z₂, …, Z_N] · [B₁; B₂; …; B_N] = Σ_{r=1}^{N} Z_r·B_r

关键观察：每张卡 r 算出 Z_r·B_r 是一个 部分和，需要把所有 N 个部分和相加才能得到完整的 Y。这个"相加"就是一次 All-Reduce。

1
2
3
4
5
6
7
8
9

列切 + 行切的配合（N = 2）：

  输入 X（每张卡都有一份完整的 X）
      │
      ├── 卡 0：Z₀ = GeLU(X·A₀)   ──→ W₀ = Z₀·B₀ ──┐
      │                                              ├── All-Reduce(sum) ──→ Y
      └── 卡 1：Z₁ = GeLU(X·A₁)   ──→ W₁ = Z₁·B₁ ──┘
                                                   ↑
                                    前向只在这里通信，且只通信一次

为什么必须"第一个按列切、第二个按行切"：

• 若两个都按列切，那么每张卡得到的是 Y 的一个列块，不需要通信——但下一层又需要完整输入，于是通信被推迟到某处，且会引出"需要拼接"的问题；
• 若两个都按行切，那 X 也必须按列切，而 X 是上一层的输出、每张卡都持有完整副本，会造成冗余计算；
• 只有"列切→行切"这一组合，能让 GeLU 这一非线性算子作用在每张卡的局部结果上而不需要通信（因为按列切得到的 Z_r 已经是完整结果的列块），并且把所有通信压缩到最后的一次归约上。

这是 Megatron-LM 全部设计的核心，值得用一个等式记住：

Y = Σ_r ( GeLU(X·A_r) · B_r )，A_r 是 A 的列块，B_r 是 B 的对应行块。

21.5.2 注意力按头切

多头注意力天然可切：H 个头之间没有计算耦合（各自独立的 Q/K/V 投影、独立的注意力、最后拼接）。因此把 H 个头分给 N 张卡，每卡 H/N 个头，就完成了一次干净的切分。

输出投影 W_O 则对应按行切（因为各卡产出的是输出维度的不同段），于是又是一次 All-Reduce。结构与 MLP 完全同构。

21.5.3 通信位置与次数

【2026-09-02 选择 Q12】 的考点就是这一段。逐项数清：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

一个 Transformer 层内（TP 度数为 N）：

  前向（2 次 All-Reduce）：
    ① 注意力输出投影之后        ← 各卡的部分和相加
    ② MLP 第二个线性层之后       ← 各卡的部分和相加

  反向（对应位置各 1 次 All-Reduce，共 2 次）：
    ③ ② 对应的反向：把激活梯度对齐（∂L/∂Z_r 需要各卡的贡献）
    ④ ① 对应的反向：同理

  合计：每层每步 4 次 All-Reduce
        （若把反向的两次记为"与权重梯度 All-Reduce 共用"，不同实现略有差异，
          但量级一定是"每层前向与反向各有若干次"）

对照真题选项：A（仅在 Loss 阶段）、C（仅在训练结束后的参数同步）、D（仅在数据加载）全部错误。正确答案是 B：每个 Transformer 层内部的前向和反向传播过程中。

21.5.4 通信量：与激活同阶

TP 的 All-Reduce 搬运的是什么？是张量 Y ∈ R^{n×d}（n 是 token 数、d 是隐藏维度），通信量约 2·n·d 个元素。

对比同一层的计算量（以 MLP 为例，2 × 2·n·d·h）：

1
2
3

每层计算量 ≈ 4·n·d·h  FLOPs
每层通信量 ≈ 2·n·d    个元素（前后向合计约 8·n·d）
计算通信比 ≈ h         ← 隐藏维度，典型 4096~8192

这个比值很健康：h 越大，每搬运一个元素能换来的计算越多。所以 TP 的单次通信量并不大。它的问题不在量，而在频率——每层每步 4 次，32 层就是 128 次，模型的每一步迭代都如此。高频率 + 高延迟的通信正是跨机 TP 掉吞吐的机制。

21.5.5 序列并行

TP 下每张卡都持有完整输入的副本，于是 LayerNorm、Dropout 这些逐元素算子在每张卡上重复计算并重复存储激活。序列并行（sequence parallelism） 沿序列维度切分这些部分，消除冗余：只在 TP 区域内沿 n 维切分，需要完整激活的算子（如注意力）前用 All-Gather 拼回来，之后再用 Reduce-Scatter 切回去。

注意这里通信量也是约 2nd 级别，与 TP 本身的 All-Reduce 同阶，因此序列并行基本是"免费"的显存优化——代价只是多两次集合通信。

21.6 流水线并行：按层切分与 bubble

21.6.1 基本形态

流水线并行（pipeline parallelism, PP）把模型的 L 层按顺序切成 p 个 stage，每个 stage 放在一张卡（或一组卡）上。第 r 个 stage 只持有 L/p 层，显存压力随之下降。

代价是顺序依赖：stage 1 必须先算完，stage 2 才能开始。为了让所有卡同时忙碌，把一个 batch 切成 m 个 micro-batch，让它们像流水线一样依次流过各 stage：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

朴素流水线（4 个 stage，4 个 micro-batch）：

        时间 →
stage1  [F1][F2][F3][F4]                        ← 前向
stage2      [F1][F2][F3][F4]
stage3          [F1][F2][F3][F4]
stage4              [F1][F2][F3][F4]
stage4                  [B4][B3][B2][B1]        ← 反向（逆序）
stage3                      [B4][B3][B2][B1]
stage2                          [B4][B3][B2][B1]
stage1                              [B4][B3][B2][B1]

        ↑ 开头和结尾这段"只有部分卡在工作"的时间就是 bubble

21.6.2 bubble 比例公式

设单个 micro-batch 通过一个 stage 的时间为 t。朴素调度下：

• 总时间 = 前向填充 (p−1)·t + 稳态 m·t + 反向排空 (p−1)·t……用更标准的口径：总步数为 2(m + p − 1)·t，其中有效工作是 2m·t，其余是 bubble：

bubble 比例 = (p − 1) / (m + p − 1)

推导思路：p 个 stage 在稳态时全部忙碌，但流水线有 p−1 个"填充槽位"和 p−1 个"排空槽位"，合计 2(p−1) 个空闲 stage-时间；总 stage-时间为 2p·m（每 micro-batch 每 stage 前后向各一次共 m·p 对）——两者之比化简即得上式。

结论：bubble 比例随 p 增大而增大，随 m 增大而减小。

算例：p = 8、m = 32 时 bubble = 7/39 ≈ 17.9%；若把 m 提到 128，bubble = 7/135 ≈ 5.2%。

21.6.3 1F1B 与虚拟流水线

1F1B（one-forward-one-backward） 调度改变了每个 stage 的执行顺序：在填充阶段结束后，每个 stage 交替执行"一个前向、一个反向"。

它的两大好处：

1. 显存峰值下降：朴素调度要先做完所有 m 个 micro-batch 的前向，才能开始反向，中途必须保存 m 份激活；1F1B 中每个设备同时"在途"的 micro-batch 数被流水线深度限制在 p 左右，因此激活显存峰值从 O(m) 降到 O(p)——这是一个与 micro-batch 数解耦的独立收益。
2. 反向可以更早开始，与后续前向重叠。

虚拟流水线（virtual pipeline / interleaved 1F1B） 让每个设备持有多个不连续的层块（如 2 个 chunk），把它们交错执行。此时每个 stage 的填充深度从 p−1 降到约 p/v（v 为虚拟 stage 数），于是

bubble 比例 ≈ (p − 1) / (m·v + p − 1)

p = 8、v = 2、m = 32 时 bubble ≈ 7/71 ≈ 9.9%——相比朴素 1F1B 的 17.9% 减半。代价是每张卡要持有 v 组层，通信次数也变成 v 倍（点对点通信本身很便宜，所以通常划算）。

21.6.4 PP 的通信特征

PP 只在相邻 stage 之间传递激活与激活梯度，是**点对点（point-to-point）**通信，通信量约 n·d 个元素/次，频率为每 micro-batch 每边界 2 次。相比 TP，它的频率低得多、且没有集合通信的同步开销。

因此 PP 是跨机扩展的首选。 唯一的代价是 bubble——bubble 期间部分卡闲置，这是买了显存和扩展性，付出了算力利用率。

21.7 ZeRO 与 FSDP：把训练状态也切开

21.7.1 训练时的显存账本

混合精度 + Adam 训练一个参数量为 P 的模型，每张卡（DP 情形下）需要存：

注意 16 字节这个数字：它是"混合精度 + Adam"的固有开销，与模型结构无关。也就是说，训练 1 个参数需要 16 字节的状态——这就是为什么训练比推理的显存需求大得多（推理 FP16 只需 2 字节/参数，见第 19 章 19.4）。

所以 7.5B 模型的标准 DP 需要 120 GB/卡，比 A100 80GB 还大——模型本身放得下（只有 15 GB），但训练状态放不下。这正是 ZeRO 要解决的问题。

21.7.2 ZeRO 的三个阶段

ZeRO（zero redundancy optimizer） 的核心思想：DP 里每张卡都存了一份完全冗余的训练状态。既然大家算出的梯度相同、更新也相同，为什么每张卡都要存全量？把它们切分到各卡上，需要时再临时取回。

为什么 ZeRO-1 只有 2.9 倍而不是 8 倍：因为优化器状态（90 GB）虽然被切分了，但参数（15 GB）与梯度（15 GB）还停留在每卡全量，合计 30 GB 是"不可压缩的地板"。当优化器状态从 90 降到 11.25 后，剩下的 30 GB 反而成了主要部分。这解释了一个常见困惑：ZeRO-1 的收益远小于 8 倍。 只有把参数也切了（ZeRO-3），才能真正逼近 N 倍。

通信代价：

这个通信量分析与 21.2.2 紧密相关：能只用 Reduce-Scatter 就绝不做 All-Reduce，因为后者是前者的两倍。基线 DP 的一次 All-Reduce 代价是 2D；ZeRO-1/2 用"Reduce-Scatter（D）+ All-Gather 更新后的参数（D）“两步替代，合计也是 2D 的两倍量级——也就是说，ZeRO 用约 2 倍的通信换来了数倍的显存节省，这是一笔明确标价的交易，不是免费的。所以能靠 ZeRO-1 解决就不上 ZeRO-3。

21.7.3 FSDP

FSDP（fully sharded data parallel） 可以理解为"ZeRO-3 的工程化实现”：参数、梯度、优化器状态全部分片，前向/反向时按需 All-Gather 某一层的参数，用完立刻释放。

1
2
3

FSDP 的一次前向：
  进入层 l  → All-Gather 该层的完整参数（临时）→ 计算 → 立即释放参数分片
  按需逐层进行，显存峰值只需容纳"当前层 + 预取的下一层"

它与 ZeRO-3 的差别主要在实现层次：ZeRO 是 DeepSpeed 库层面的优化，FSDP 是 PyTorch 原生集成的实现，并支持嵌套 FSDP / HSDP（hybrid sharded data parallel）——把 N 张卡分成若干组，组内分片、组间复制，从而把 All-Gather 的通信限制在更小、更快的域内。这又一次体现了 21.3 的规则：让高频通信留在快域里。

21.8 专家并行与 MoE

21.8.1 稀疏激活

MoE（mixture of experts） 把 MLP 层替换成 E 个并行的"专家"，每个 token 只被路由到 top-k 个专家（常见 k = 1 或 2）。

1
2
3
4
5
6
7
8

MoE 层：
  输入 token x
      │
      ├── Router（一个小线性层 + softmax）→ 选出 top-k 个专家
      │
      ├── 只把 x 发给被选中的 k 个专家做 FFN
      │   其余 E−k 个专家这一轮完全不参与
      └── 按 gating 权重加权求和

收益：总参数量可以做到 E 倍，但每个 token 的实际计算量只增加 k 倍。这就是稀疏激活的价值——参数量（决定知识容量）与计算量（决定成本）解耦了。

21.8.2 All-to-All 通信

专家并行（expert parallelism, EP）把不同专家放在不同设备上。于是每个 token 必须被发送到它所选专家所在的设备上，算完再送回来——这是一次 All-to-All：

1
2
3
4
5

EP 的一次 All-to-All（2 个阶段）：
  ① Dispatch：每张卡把自己的 token 按目标专家分桶，各桶分别发给对应的卡
  ② Combine ：各卡把专家算完的结果按原路送回

  通信量 ≈ 2 × k × n × d 个元素（k = top-k，n = token 数，d = 隐藏维度）

注意这里出现了 n（token 数）：EP 的通信量与 token 数成正比，与 TP 的 n·d 同阶但乘上了 k 与路由的不均匀性。而且 All-to-All 是每层都做——所以 EP 的通信频率与 TP 同级，必须放在 NVLink 域内或用高带宽 IB。

21.8.3 负载均衡与容量因子

MoE 最棘手的工程问题是负载不均：router 可能把所有 token 都送给少数几个专家，导致那几张卡算不完、其余卡闲置。两种治理手段：

• 负载均衡损失（auxiliary load-balancing loss）：在训练目标中加一项，惩罚"专家使用率方差过大"，用梯度把流量推均匀。这是训练层面的手段。
• 容量因子（capacity factor, CF）：工程层面的手段——为每个专家设定一个容量上限 capacity = CF × (n·k / E)，超出上限的 token 被直接丢弃（走残差旁路，不经过专家）。
  • CF = 1.0 表示容量恰好等于平均负载，此时任何不均衡都会导致丢 token；
  • CF = 1.25 是常见折中：给 25% 的余量，既能吸收一定波动，又不会浪费太多算力；
  • CF 越大越不容易丢 token，但显存与算力浪费也越大。

21.9 3D/4D 并行的组合选型

真实的大规模训练从来不是"选一种并行"，而是多种并行叠加，常见的表述如 “TP=8, PP=8, DP=64”（共 4096 卡）。

21.9.1 决策流程

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

开始：模型能否在单卡上完成一次训练迭代？
  │
  ├── 能 ⟹ DP（+ZeRO-1/2 降低显存），不需要 TP/PP
  │
  └── 不能
       │
       ├── 是"参数放不下"（单层就能放下）？
       │      └── 用 PP（按层切）+ DP。PP 跨机安全。
       │
       ├── 是"单层放不下"（比如注意力或 MLP 的权重矩阵过大）？
       │      └── 必须用 TP，且 TP 度数 ≤ 单机卡数（≤ NVLink 域大小）
       │
       └── 是"激活/序列太长导致放不下"？
              └── 用 TP + 序列并行 + 激活重算（gradient checkpointing）

三条经验规则：

1. TP 度数不超过单机的 GPU 数（通常 8，GB200 NVL72 这类系统可更大）。这是 21.3 的直接推论。
2. PP 度数 ≤ 层数 / 每卡最少层数，且要让 m ≥ 4p 以把 bubble 压到可接受范围（21.6.2 的公式）。
3. DP 度数吃掉剩下的所有卡，因为 DP 最容易扩展、通信频率最低。

21.9.2 全局 batch size 的约束

全局 batch = micro_batch × 梯度累积步数 × DP 度数

三条约束同时成立：

• micro_batch 受单卡显存限制（尤其受激活占用限制，与序列长度平方相关）；
• 梯度累积步数用于把全局 batch 补到目标值，同时降低 DP 通信频率（21.4.2）；
• DP 度数 是硬件与并行策略给定的。

因此实践中常常是：先定全局 batch（由收敛性决定），再由显存定 micro_batch，最后用梯度累积补足差额。 这也解释了为什么改并行配置常常需要连带改学习率——全局 batch 变了，学习率的缩放规律（线性缩放或平方根缩放）也要跟着调整。

21.10 通信开销的度量与诊断

有了上面的框架，诊断"训练慢"就有了清晰的路径：

通信占比（communication ratio） 是最核心的单一指标：

通信占比 = 纯通信时间 / 每步总时间

若该值小于约 10%，说明重叠做得好，可以继续加卡；若超过 30%，则加卡的收益已经很低，应当先优化通信（换更快的域、调整 bucket、减小 TP 度数）。

计算通信重叠是把通信占比压下去的主要手段，DDP 的 bucket（21.4.2）是最经典的例子：把通信拆成小块，与后续计算并行发起。TP 的 All-Reduce 也可以与后续层的计算重叠（但重叠窗口比 DP 小，因为每层之间耦合更紧）。

本章小结

自测题

Q1（推导）N = 8 张卡做数据并行，模型参数量 P = 7×10⁹（FP16 梯度）。求每次梯度同步的 All-Reduce 通信量，并估算在 100 Gb/s 以太网（约 12.5 GB/s）与 NVLink（约 450 GB/s 单向）上分别需要多久。

1
2
3
4

以太网（12.5 GB/s）：24.5 GB / 12.5 GB/s ≈ 1.96 s
                     即使按乐观的单向 12.25 GB 算，也要 ≈ 0.98 s
NVLink（450 GB/s） ：24.5 GB / 450 GB/s ≈ 54 ms
                     单向口径下 ≈ 27 ms

Q2（推导）写出 Megatron 风格 TP 中 MLP 的切分方案，并证明 Y = GeLU(X·A)·B 只需一次 All-Reduce。若 A ∈ R^{4096×16384}、X ∈ R^{n×4096}、n = 2048，在 N = 8 时估算每张卡的显存占用节省比例。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

把 A 按列切成 N 份：A = [A₁, A₂, …, A_N]，A_r ∈ R^{d × h/N}
把 B 按行切成 N 份：B = [B₁; B₂; …; B_N]，B_r ∈ R^{h/N × d}

卡 r 计算：Z_r = GeLU(X · A_r) ∈ R^{n × h/N}
           W_r = Z_r · B_r     ∈ R^{n × d}

由 GeLU 逐元素可得：
  GeLU(X·A) = GeLU(X·[A₁,…,A_N]) = [GeLU(X·A₁), …, GeLU(X·A_N)] = [Z₁,…,Z_N]
于是
  Y = GeLU(X·A)·B = [Z₁,…,Z_N]·[B₁;…;B_N] = Σ_r Z_r·B_r
                                                └── 各卡的部分和

Q3（计算）一个 32 层模型用 PP 切成 p = 8 个 stage，micro-batch 数 m = 32。(a) 朴素 1F1B 的 bubble 比例是多少？(b) 用 v = 4 的虚拟流水线后是多少？(c) 要把 bubble 压到 5% 以下，m 需要多大（v = 1）？

Q4（计算）P = 7.5×10⁹、N = 8 的混合精度 Adam 训练。分别估算基线 DP、ZeRO-1、ZeRO-2、ZeRO-3 的每卡显存，并解释为什么 ZeRO-1 只有约 3 倍收益而不是 8 倍。

1
2
3
4

基线 DP ：15 + 15 + 90            = 120.0 GB
ZeRO-1  ：15 + 15 + 90/8 = 11.25  =  41.25 GB   （2.9×）
ZeRO-2  ：15 + 15/8 + 11.25       =  28.1 GB    （4.3×）
ZeRO-3  ：15/8 + 1.875 + 11.25    =  13.1 GB    （9.2×）

Q5（辨析）某团队把 TP 度数设为 16，跨了两台 8 卡机器（机间 InfiniBand）。他们发现扩展后单步耗时反而比 TP=8 更长。请解释原因，并给出两种改法。

与真题的关联

延伸阅读

• Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019 —— 列切/行切方案与每层 All-Reduce 的原始出处。
• Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC 2020 —— 三阶段切分与通信量分析。
• Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021 —— 1F1B 调度、虚拟流水线与 3D 并行的组合选型。
• Huang et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019 —— bubble 比例公式的来源。
• Fedus et al., Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, JMLR 2022 —— 容量因子、负载均衡损失与丢 token 问题。
• Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, MLSys 2023 —— 序列并行与激活重算的完整分析。
• 第 19 章（推理系统与性能工程）提供"算力 / 带宽 / 算术强度"的量化工具；第 20 章（量化与压缩）提供降低单卡显存占用从而减少并行度需求的手段。