第 16 讲 · 分布式训练（三）：模型并行——流水线并行与张量并行 · Distributed Training III: Model Parallelism

对应 Maps to：CMU 11-868 / 11-968 Large Language Model Systems（Spring 2026，Lei Li），Lecture 16 · Mon 3/16 · Distributed Model Training III（slides：llmsys-16-model-parallel.pdf，39 页）。 阅读材料：Huang et al. (2019) GPipe；Narayanan et al. (2019) PipeDream；Shoeybi et al. (2019) Megatron-LM；Narayanan et al. (2021) Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM（SC）；Chen et al. (2016) Training Deep Nets with Sublinear Memory Cost；Korthikanti et al. (2022) Reducing Activation Recomputation in Large Transformer Models。 依据说明：slides 覆盖：模型规模的增长与模型并行的动机、按层切 / 按张量切、朴素的流水线与它的三个缺点、GPipe 的微批次与气泡 (K − 1)/(M + K − 1)、吞吐表、激活重计算（O(n)、O(1)、√n 三种）、GPipe 的在途激活、1F1B（PipeDream-Flush）、交错式调度（Megatron）、shallowspeed 的 GPipe 代码、torch.distributed.pipelining 的用法、张量并行（矩阵按列 / 按行切、MLP、注意力、嵌入与输出层、LayerNorm 复制）、流水线 + 张量 + 数据并行的组合与 Megatron 的两条结论。HW5 Problem 2 的 clock_cycles 与 Pipe 取自作业说明。f / g 算子、按词表切分的交叉熵的具体做法、序列并行与激活显存公式、气泡与 d 的关系取自上面几篇论文。本讲义的调度模拟器、线程版流水线运行时与本机计时、手写的重计算、多进程（gloo）张量并行层、3D 并行的代价模型，都是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture16/（6 个脚本 + pipeline.py、tensor_parallel.py；每个都在 5 秒以内） 预计阅读时间：3 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 说清楚为什么要模型并行、朴素地”按层切”为什么只有 1/K 的利用率，以及按层切、按矩阵切、按数据切各自的通信特点。
2. 写出 HW5 的 clock_cycles，推导 GPipe 的气泡比例 (K − 1)/(M + K − 1)，并解释微批为什么不是越多越好。
3. 实现一个按波次提交的流水线前向（HW5 的 Pipe），知道它为什么和顺序执行”数学上相同、浮点上只差舍入”，避开线程里的几个坑。
4. 比较 GPipe、1F1B、交错式调度的时间与激活显存；用激活重计算以计算换显存，算出 √n 检查点的代价。
5. 推导 Megatron 张量并行的切法（MLP 先列后行、注意力按头、输出层按词表），写出 f / g 两个通信算子，算出每层的通信量。
6. 用 Megatron 的两条结论和一个代价模型，为给定的模型和集群选 (t, p, d)。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  为什么要模型并行：按层切；朴素的流水线；三种切法的通信强度                      ← 代码 01
§2  GPipe：微批次、时钟周期调度表、气泡；微批的个数                                ← 代码 02
§3  实现一条流水线：HW5 的 Pipe、正确性、几个坑、⏱ 本机加速                        ← 代码 03
§4  省显存：1F1B、交错式调度、激活重计算                                          ← 代码 04
§5  张量并行：列切分与行切分、MLP、注意力、输出层、f / g、序列并行                  ← 代码 05
§6  3D 并行：张量 × 流水线 × 数据                                                ← 代码 06
§7 小结   §8 常见误区   §9 练习   §10 延伸阅读   §11 下一讲预告

1. 为什么要模型并行 · Why model parallelism

1.1 一块卡放不下

第 14、15 讲的数据并行有一个前提：每块卡都放得下整个模型。slides 第 3 页的曲线说明这个前提早就不成立了：GPT-3 175B、LLaMA-3.1 405B、DeepSeek-V3 671B。第 14 讲算过训练状态（混合精度 Adam 每参数 16 字节）：175B 要 2.8 TB，而一块卡只有 80 GB。

那就把模型的计算切开，放到几块卡上（slides 第 4 页）。有两种基本切法：

• 按层切（layer-wise）：卡 0 放第 0 层，卡 1 放第 1 层……前向时激活从卡 0 流到卡 3，反向时梯度从卡 3 流回卡 0。这叫流水线并行（pipeline parallelism）。
• 按张量切（tensor-wise）：每一层里的矩阵都切开，每块卡算矩阵乘法的一部分。这叫张量并行（tensor parallelism）。

slides 还提到第三种——专家并行（MoE 的专家放在不同卡上），留到第 17 讲。

01 号脚本第 1 部分按层切 GPT-3（96 层、h = 12288）：

只做推理，切 8 段就够；训练要切到 48 段——还没算激活。按层切很自然：Transformer 是一摞相同的层，切开之后每块卡上是一个完整的小模型，段与段之间只传一个激活张量。唯一的不均衡是第一段多了词嵌入（相当于 0.35 层）。

1.2 朴素的模型并行：任何时刻只有一块卡在干活

最直接的做法（slides 第 5 页）：一个批先在卡 0 上算完前 L/K 层，把输出（nccl send/recv）交给卡 1……最后一块卡算出损失，反向再一路传回来。01 号脚本第 2 部分画出 K = 4 的时间线：

设备 0 |F0F0····································B0B0B0B0|
设备 1 |····F0F0························B0B0B0B0········|
设备 2 |········F0F0············B0B0B0B0················|
设备 3 |············F0F0B0B0B0B0························|

利用率正好 1/K = 25%——前向是一层接一层的，卡 1 必须等卡 0 算完。slides 第 8 页列了朴素做法的三个问题：

1. 利用率低：任何时刻只有一块卡在工作（上图）。
2. 计算与通信不重叠：段间传输时两块卡都在等（加上 0.5 个单位的传输，利用率降到 20%）。
3. 显存压力大：卡 0 最早做完前向，却要一直存着整个批的激活，直到最后才轮到它做反向。

GPipe 的微批次同时缓解这三件事（§2）。

1.3 三种切法的通信强度

切法不同，要传的东西完全不同。01 号脚本第 3 部分用一个好懂的指标比较：每块卡每秒要收发多少字节，才跟得上它的计算（GPT-3，一条 2048 token 的序列，按 A100 峰值算）：

（对照：机器内部 NVLink 每块卡约 300 GB/s，跨机器 InfiniBand 每块卡约 25 GB/s。）

• 流水线并行通信极少：只在段的边界上传一个 s × h 的激活，而每段要算一整串层。它可以跨机器。
• 张量并行每一层都要 all-reduce 一个完整的激活，通信强度高出两个数量级：NVLink 撑得住，InfiniBand 撑不住——这是 §5、§6 的主线。
• 数据并行每步要 all-reduce 和模型一样大的梯度，但每步只一次、可以和反向重叠（第 15 讲），每卡的批越大越划算。它的问题是每块卡要放下整个模型。

所以真正的大模型训练是三者的组合（§6）。

2. GPipe：微批次 · Micro-batching

2.1 想法

GPipe（Huang et al. 2019，slides 第 9–11 页）的想法很简单：把一个批（mini-batch）切成 M 个微批（micro-batch）。卡 0 算完微批 0 的前向，把输出交给卡 1，然后马上开始算微批 1——卡 0 和卡 1 同时在干活，就像工厂的流水线。所有微批的前向做完之后再做反向，梯度在各微批之间累加，最后统一更新参数（slides 第 18 页 shallowspeed 的 minibatch_steps）。

为什么这样做在数学上是对的？因为参数在一个批之内不变，整批的梯度 = 各微批梯度之和（损失按样本平均时再除以 M）。微批只改变了计算的时间安排，不改变计算本身（§3.2 会验证）。

2.2 时钟周期调度表（HW5 Problem 2.1）

假设每个 (微批, 阶段) 的前向都花一个时钟周期。微批 i 在阶段 j 上的前向要等两件事：它在阶段 j − 1 上算完（数据依赖），阶段 j 算完微批 i − 1（一块卡一次只做一件事）。所以它最早在时刻 i + j 开始。第 k 个时钟周期（一”波”）里能做的，就是所有满足 i + j = k 的 (i, j)：

clock_cycles(3, 3)：
  k = 0: [(0, 0)]
  k = 1: [(1, 0), (0, 1)]
  k = 2: [(2, 0), (1, 1), (0, 2)]
  k = 3: [(2, 1), (1, 2)]
  k = 4: [(2, 2)]

这正是 HW5 给出的例子。实现（pipeline.clock_cycles）只有几行：对每个 k，阶段 j 的范围是 [max(0, k − M + 1), min(k, K − 1)]，微批 i = k − j。作业要求的几个边界：

• 一共 M + K − 1 波；每一波按阶段 j 从小到大排列；pair 写成 (微批, 阶段)——注意 slides 的图里 F_{a,b} 是”阶段 a、微批 b”，顺序相反。
• M = 0 时什么都不产生；M < 0 或 K < 1 抛 ValueError。它是生成器，调用时不会报错，第一次迭代时才报错（02 号脚本第 1 部分演示了这一点，HW5 明确允许）。

2.3 气泡

前 K − 1 波，流水线在”灌满”，后面的阶段空着；最后 K − 1 波在”排空”，前面的阶段空着。这些空闲叫气泡（bubble）。把反向也画进来（反向 = 2 × 前向），02 号脚本第 2 部分：

K = 4，M = 8：一步 33 个单位，气泡 27.3%
  设备 0 |F0F1F2F3F4F5F6F7··················B7B7B6B6B5B5B4B4B3B3B2B2B1B1B0B0|
  设备 1 |··F0F1F2F3F4F5F6F7············B7B7B6B6B5B5B4B4B3B3B2B2B1B1B0B0····|
  设备 2 |····F0F1F2F3F4F5F6F7······B7B7B6B6B5B5B4B4B3B3B2B2B1B1B0B0········|
  设备 3 |······F0F1F2F3F4F5F6F7B7B7B6B6B5B5B4B4B3B3B2B2B1B1B0B0············|

左上和右下两个三角形就是气泡。推导：每块卡要做 M 个前向、M 个反向，忙 M·(t_f + t_b)；一步总长 (M + K − 1)·(t_f + t_b)。所以

$\text{气泡比例} = 1 - \frac{M}{M + K - 1} = \frac{K - 1}{M + K - 1}$

（slides 第 11 页；按”相对有效计算”的口径则是 (K − 1)/M。）模拟器与公式的最大差 5.6e-17：

GPipe 论文说 M ≥ 4K 时气泡开销可以忽略——按公式，M = 4K 时气泡仍有 16–18%，并不小。论文的说法还依赖另一件事：反向里的重计算（§4.3）不需要等后面的梯度，可以提前开始、填进气泡。只看公式，要把气泡压到 5% 以下，M 要大于 19(K − 1)。

吞吐（02 号脚本第 3 部分，相对”整个模型在一块卡上”）：M = 1 时多少块卡都只有 1 倍；K = 8、M = 32 时 6.56 倍（公式 K·M/(M + K − 1)）。段间传输等于半个前向时，M = 1 反而比单卡慢（0.77–0.86 倍），M = 32 时只少一点——传输和别的微批的计算重叠了。这就是 slides 第 12 页的吞吐表的规律：微批越多越接近线性加速，在没有高速互连的机器上也能用。

2.4 微批不是越多越好

气泡随 M 变小，那为什么不把 M 设成批大小？因为每个微批的计算都有固定开销：内核发射、通信启动、小矩阵用不满计算单元。02 号脚本第 4 部分在本机实测一段 MLP 处理 b 行的耗时（⏱），拟合出 t(b) ≈ α + β·b（α 约 60–80 μs）：

（⏱ 本机数字，每次运行都会不同。）一步的时间先降后升。把一步写成 3·(M + K − 1)(α + βN/M)，对 M 求最小，得到 M* = √((K − 1)βN/α)，这台机器上约 15–17，和扫描的结果一致。GPU 上 α 更大，所以每个微批至少要有几百到上千个 token。slides 第 9 页说”批由显存决定，尽量填满；微批可以更小”——但也不能太小。

3. 实现一条流水线 · Implementing a pipeline

3.1 HW5 的 Pipe

HW5 Problem 2.2 给每个阶段一个工作线程（StageWorkers：submit(stage, fn) 把函数交给那个阶段的线程，receive(stage) 等它的结果），要求 Pipe.forward：

batches = list(torch.split(inputs, self.split_size))
for wave in clock_cycles(len(batches), len(self.stages)):
    for i, j in wave:                          # 先把这一波全部提交
        self.workers.submit(j, lambda m=self.stages[j], x=batches[i]: m(x))
    for i, j in wave:                          # 再按同样的顺序逐个接收
        batches[i] = self.workers.receive(j)
return torch.cat(batches)

03 号脚本第 1 部分打印了 M = K = 3 时的调用顺序：

波 2：submit(0)  submit(1)  submit(2)  receive(0)  receive(1)  receive(2)

两个要点：一波之内先全部 submit 再 receive，同一波的几个阶段才能同时算（submit 一个就 receive 一个会退化成顺序执行）；一波收齐才开始下一波，因为下一波的 (i, j + 1) 要用这一波 (i, j) 的输出。HW5 的评分器记录每次调用，顺序必须完全一致。在 GPU 上，跨设备的 .to(device) 也要放进提交的函数里——放在发射线程里会让各阶段串行地等拷贝，HW5 提示这样拿不到 1.10× 的加速。

3.2 和顺序执行一样吗？

03 号脚本第 2 部分把同样的 3 段模型分别用顺序执行和 Pipe 跑（批 10 行，split_size = 4 → 微批 4、4、2 行，最后一个更小是允许的）：

数学上相同，浮点上只差舍入。原因有两个：

• 参数的梯度是对批里所有样本求和。流水线是”每个微批算一个部分和，再由 autograd 累加”，求和的顺序变了——和第 14 讲数据并行”不逐位相同”是同一个道理。
• 连输出都差一点：每一行的数学运算没变，但矩阵乘法库会按矩阵的形状选择分块和累加顺序，4 行和 10 行的矩阵乘法，同一行的结果可能差在最后一位。

所以 HW5 的反向评分用带容差的比较，对比的是”同样的几段依次执行、不做流水线”。注意不要 detach 中间结果：反向传播要沿着它们回到前面的阶段。

3.3 三个坑

03 号脚本第 3 部分：

1. torch.no_grad() 在工作线程里失效。PyTorch 的梯度开关是每个线程各一份的（thread-local），调用方的 no_grad 管不到工作线程。我们的 StageWorkers 在 submit 时记下开关、在工作线程里设成一样的；HW5 的评分也检查 no_grad 下没有建图。
2. 异常要原样抛出。阶段抛出的异常在工作线程里被捕获，receive 时重新抛出，类型和信息不变（不要包装成别的异常，也不要吞掉）。
3. lambda 的晚绑定。[lambda: i for i in range(3)] 执行后是 [2, 2, 2]：闭包记住的是变量，不是当时的值。交给线程的函数稍后才执行，所以要写成 lambda m=module, x=batch: m(x)，用默认参数把当时的值绑进去。

3.4 ⏱ 本机的流水线加速

HW5 Problem 2.3 在两块 V100 上要求 ≥ 1.10× 的加速（2 段、每段 4 个 MLP 块、8192 行切成 8 个 1024 行的微批）。03 号脚本第 4 部分在本机 CPU 上跑缩小版（每个阶段一个线程，PyTorch 的算子执行时释放 GIL，两个线程能在两个核上同时算）：

前向在 8 个微批时快约 1.4–1.5 倍（两次运行分别是 1.52 和 1.37；上限 2·8/9 = 1.78）。加上反向就几乎没有加速：HW5 的 Pipe 只流水线化了前向，反向交给 autograd。GPU 上 autograd 给每块卡一个线程，不同卡上的反向能部分重叠；CPU 上 autograd 只有一个线程，反向完全顺序。要让反向也按时间表执行，就得自己调度每个微批的反向——这正是 torch.distributed.pipelining 做的事（slides 第 21–24 页）：

1. 把模型切成 PipelineStage。可以手动切：在 meta 设备上构造整个模型，每个 stage 删掉不属于自己的层——slides 的 Transformer 用 ModuleDict 存层、forward 里允许 tok_embeddings/norm/output 为 None，就是为了能这样切。也可以自动切。
2. 选一个调度：ScheduleGPipe(stage, n_microbatches)（也有 1F1B、交错式）。rank 0 调用 schedule.step(x)，批会被自动切成微批；其他 rank 调用 schedule.step()。

4. 省显存：1F1B、交错式调度、激活重计算 · Saving memory

4.1 GPipe 的问题：M 个微批同时在途

GPipe 先做完所有前向再做反向，于是每块卡在反向开始之前要同时存着 M 个微批的激活（slides 第 15 页）。M 越大气泡越小，激活却越多——两个目标打架。

1F1B（PipeDream-Flush，Narayanan et al. 2019，slides 第 16 页）调整顺序：卡 d 先做 K − d − 1 个”预热”前向，然后一个前向、一个反向交替，最后把剩下的反向做完。反向越早开始，越早能释放那个微批的激活。04 号脚本第 1 部分：

1F1B（K = 4，M = 8）：一步 33 个单位，气泡 27.3%，每块卡在途微批的峰值 [4, 3, 2, 1]
  设备 0 |F0F1F2F3············B0B0F4B1B1F5B2B2F6B3B3F7B4B4··B5B5··B6B6··B7B7|
  设备 1 |··F0F1F2········B0B0F3B1B1F4B2B2F5B3B3F6B4B4F7B5B5··B6B6··B7B7····|
  设备 2 |····F0F1····B0B0F2B1B1F3B2B2F4B3B3F5B4B4F6B5B5F7B6B6··B7B7········|
  设备 3 |······F0B0B0F1B1B1F2B2B2F3B3B3F4B4B4F5B5B5F6B6B6F7B7B7············|

• 一步的时间、气泡和 GPipe 完全一样（33 个单位）：1F1B 只是把反向往前挪，没有减少空闲。
• 在途激活从 8 个降到 4 个（卡 d 是 K − d 个），与 slides 的”最多 4 个 vs 8 个”一致。关键是这个上限只和 K 有关、与 M 无关——于是可以放心地增大 M 来压气泡。这是 1F1B 成为默认调度的原因。

一个细节：段间传输不为 0 时（04 号脚本第 3 部分下半张表），1F1B 的气泡反而比 GPipe 大（M = 32、传输 0.2 时 16.8% vs 9.6%）。GPipe 一口气做完所有前向，传输被后面的微批遮住；1F1B 的稳定阶段里每个反向都要等后一块卡传回梯度，传输延迟每一步都落在关键路径上。真实系统用异步的 isend / irecv 尽量把传输和计算重叠。

4.2 交错式调度

Megatron-LM（Narayanan et al. 2021，slides 第 17 页）再进一步：每块卡不再持有连续的一段层，而是持有 v 段不相邻的层（“模型块”）。比如 16 层、4 块卡、v = 2：卡 1 持有第 1–2 层和第 9–10 层。每个微批在设备之间转 v 圈，但每段只有原来的 1/v 层，流水线的灌满和排空也只要 1/v 的时间。04 号脚本第 2 部分（K = 4、M = 8、v = 2）：

气泡时间正好缩小 v 倍（Megatron：气泡 = (1/v)·(K − 1)/M）。代价有三：传输次数约 v 倍；预热阶段做的前向更多，在途激活从 4 个增加到 5.5 个；还要求 M 是 K 的整数倍。传输有代价时（传输 0.2 个单位），v = 4 在每个 M 下都比 v = 2 差——v 怎么选要看互连有多快。

4.3 激活重计算

另一条省激活的路是激活重计算（activation recomputation / gradient checkpointing，Chen et al. 2016，slides 第 13–14 页）：前向时不存所有中间结果，反向需要时再重新算出来。反向第 l 层需要它的输入 h_l 和输出 h_{l+1}，三种做法：

第三种：前向时只存 h_0、h_k、h_2k……；反向处理某一段时，从这段的检查点重新前向一遍、临时存下这一段的 k 个中间结果，算完就丢。存的个数 ≈ n/k + k，k = √n 时最小（≈ 2√n）。04 号脚本第 4 部分用手写的一串 tanh 层验证（n = 64）：

梯度逐位相同：重计算用同样的输入、同样的运算再算一遍。它只是拿计算换显存——√n 检查点多做约一次前向，约占训练计算量的 1/3。

放进流水线（04 号脚本第 5 部分）：GPipe 打开重计算后，每块卡只存每个微批进入这一段时的输入（一共 N 个样本），反向时再重算一个微批在这一段的 L/K 层。峰值激活从 O(N × L) 降到 O(N + L/K × N/M)（GPipe 论文）。批 256、48 层、K = 4、M = 32 时省了 8.7 倍；M = 1 时反而多存一点（除了重算的整段激活还要存输入）——重计算要和微批配合。

5. 张量并行 · Tensor parallelism

5.1 矩阵乘法的两种切法

流水线按层切；张量并行把一层里面的矩阵切开（slides 第 27–28 页）。Y = X·A 有两种切法：

• 按列切 A = [A₁, A₂]：卡 i 算 X·A_i，拿到 Y 的一部分列（完整的值）。要得到整个 Y 就 all-gather（拼起来）。
• 按行切 A = [A₁; A₂]，同时把 X 按列切成 [X₁, X₂]：卡 i 算 X_i·A_i，拿到整个 Y 的一个部分和。要得到 Y 就 all-reduce（加起来）。

两种都精确（05 号脚本第 1 部分，与 X·A 差 8.9e-16）。区别在于后面接非线性函数时：GeLU(X₁A₁ + X₂A₂) ≠ GeLU(X₁A₁) + GeLU(X₂A₂)（差 3.3）。按行切的话，GeLU 之前必须先 all-reduce；按列切则每块卡各自做 GeLU 就行，因为 GeLU 是逐元素的，不跨列。

5.2 MLP：先列后行

Megatron-LM（Shoeybi et al. 2019，slides 第 29 页）的 MLP：Z = GeLU(X·A)·B。

1. A 按列切：卡 i 算 Y_i = GeLU(X·A_i)，不需要通信。
2. B 按行切：Y 已经是按列分好的 [Y₁, Y₂]，正好是按行切 B 需要的输入。卡 i 算 Z_i = Y_i·B_i（部分和）。
3. 一次 all-reduce：Z = Z₁ + Z₂。

整个 MLP 前向只有最后那一次 all-reduce。05 号脚本第 2 部分在 t = 2、4、8 上验证：与不切分差 2e-15 量级，每块卡只存 1/t 的权重、只做 1/t 的乘法；没有 all-reduce 时卡 0 手里的部分和与结果差约 2。

5.3 注意力：按头切

多头注意力天然可以按头切（slides 第 30 页）：每个头的计算互不相关。Wq、Wk、Wv 按列切——每一份正好是若干个完整的头；每块卡独立算自己那几个头的注意力；输出投影 Wo 按行切，接住按头分好的输出，得到部分和；最后一次 all-reduce。头数必须能被 t 整除——这也是 t 通常不超过 8 的原因之一。

5.4 f 和 g：通信算子与它们的反向

把通信写成两个 autograd 算子（Megatron 论文图 3），整个层就清楚了：

为什么 f 的反向要 all-reduce？每块卡只用自己那份权重，算出的只是”对 X 的梯度”的一部分，加起来才完整。f 和 g 互为共轭。一个 Pre-LN Transformer 层就是：

x = x + g(Attn_rank(f(LN₁(x))))
x = x + g(MLP_rank(f(LN₂(x))))

tensor_parallel.py 用 torch.autograd.Function 写了这两个算子（十来行），05 号脚本第 3 部分在本机真正起 2 个、4 个进程（gloo，float64）跑一个张量并行层的前向和反向，与单进程对比：

每层前向 2 次、反向 2 次 all-reduce，每次一个完整的激活（序列长 × h），与 t 无关。LayerNorm、dropout、残差在每块卡上各做一份（slides 第 32 页 “Duplicate across GPUs”）：γ 的梯度在每个进程上算出来都一样，所以各自更新也保持一致。行切分的矩阵的偏置不能切，只能加一次——我们干脆不用偏置；Megatron 的做法是 all-reduce 之后再加。

5.5 嵌入与输出层

词表很大（GPT-2 是 50257，LLaMA-3 是 128256），嵌入矩阵也按词表切（slides 第 31 页）：

• 输入嵌入：卡 i 只有词号落在自己那一段的嵌入，其余位置填 0，然后 all-reduce。
• 输出层：按词表切之后，每块卡得到一部分 logits。如果先 all-gather 出完整的 logits 再算交叉熵，通信量是 b·s·V。Megatron 把交叉熵融合进来：交叉熵 = logsumexp(z) − z[目标]，logsumexp 可以分块算——all-reduce 各块的最大值（数值稳定）、all-reduce 各块的 Σexp、目标词的 logit 只在一块卡上，再 all-reduce 一次。每个 token 只交换 3 个标量。

05 号脚本第 4 部分：结果与 F.cross_entropy 差 1.8e-15；t = 8、一个微批 8 × 2048 个 token 时，all-gather logits 要 1.44 GB（GPT-2 词表）或 3.68 GB（LLaMA-3），融合的交叉熵只要 0.34 MB——相差 4188 倍和 10688 倍。这和第 10 讲的”融合嵌入与损失”是同一个思想：别把大张量拼出来，把后面的计算挪到数据所在的地方。

5.6 通信与计算之比；序列并行

一层的计算（前向 + 反向）约 72·b·s·h²/t FLOPs；通信是 4 次 all-reduce，每次 2(t − 1)/t · b·s·h·2 字节。两者都正比于 b·s，所以通信 / 计算之比与批大小、序列长度无关，只和 t、h、带宽有关。05 号脚本第 5 部分（GPT-3，A100）：

Megatron 的张量并行通信在关键路径上（下一层要等这一层 all-reduce 完），不能像 DDP 那样和计算重叠。所以 t ≤ 8，只在一台机器内部用（slides 第 34 页的结论 #1）。

张量并行也省激活，但不彻底。Korthikanti et al. (2022) 的公式（16 位激活，每层）：

那个”10·sbh”来自 LayerNorm、dropout 和残差：它们在每块卡上各存一份完整的。序列并行（sequence parallelism）把这些部分沿序列维切开：原来的一次 all-reduce 拆成 reduce-scatter（进入 LayerNorm 之前，每块卡拿到 1/t 的序列）+ all-gather（进入下一个张量并行块之前），通信量不变（第 14 讲：all-reduce = reduce-scatter + all-gather），激活全部除以 t。选择性重计算只重算 s² 大小的注意力分数和 softmax 那一部分——显存省得多，计算只多几个百分点。

6. 3D 并行 · Putting it together

6.1 组合

N 块卡 = t（张量并行）× p（流水线并行）× d（数据并行）（slides 第 33–35 页）。典型的摆法：

• 张量并行组放在一台机器内部（8 块卡，NVLink）；
• 流水线的各段放在不同的机器上（点对点，通信少）；
• 最外层是数据并行：t·p 块卡组成”一份模型”，d 份副本之间 all-reduce 梯度。

Megatron 的两条结论：

1. 张量并行用到一台机器的卡数为止，再用流水线并行跨机器扩展模型（§5.6 的通信 / 计算之比）。
2. 模型并行的规模 t·p 只要大到放得下模型就够了，剩下的卡用数据并行。

6.2 一个代价模型

06 号脚本把每步的时间拆成：计算（Megatron 的公式 96·B·s·L·h²·(1 + s/(6h) + V/(16·L·h))，含整层重计算）、气泡 (p − 1)/m、张量并行的 all-reduce（每层每个微批 6 次：前向 2、重计算 2、反向 2）、流水线的点对点传输、数据并行的梯度 all-reduce；显存 = 训练状态 16 字节 × 参数 / (t·p) + 激活。对论文表 1 的万亿参数模型（128 层、h = 25600、3072 块 A100、批 3072）：

• 模型明确选出 t = 8：t = 16 时张量并行跨出一台机器，all-reduce 走 IB，通信时间暴涨；t 小于 8 时每块卡上的模型更大，要用更多段才放得下，或者微批更少，气泡变大。
• p 在 32–128 之间差别很小；p = 32 时显存 78 GB 已经贴着上限，真实系统还有临时缓冲和碎片（第 11、18 讲），p = 64 留出了余量。论文用的正是 t = 8、p = 64，实测每块卡 163 TFLOP/s（峰值的 52%）、合计 502 PFLOP/s。模型的 197 TFLOP/s 偏高，因为它假设矩阵乘法跑到峰值的 75%，也没算非矩阵乘法的时间。
• 显存的大头是训练状态：t·p = 512 时每卡约 31 GB。这也是后来 Megatron 引入”分布式优化器”（相当于 ZeRO-1）的原因——第 18 讲。

6.3 为什么 t·p 够用就好

卡数 n 和 t 固定时，p = n/(t·d)，每条流水线的微批 m = B/d，于是

$\text{气泡比例} = \frac{p - 1}{m} = \frac{n/t - d}{B}$

数据并行的 d 越大，气泡越小（Megatron 论文 §3.3）。06 号脚本第 2 部分（GPT-3、1024 块卡、t = 8）：放得下的三种配置，气泡从 7.3% 到 8.1%；数据并行的梯度 all-reduce 随 d 增加，两者在这个模型里几乎抵消，一步只差 0.3% 以内。论文在真实集群上测得的差距更明显（流水线段越多吞吐越低）：段越多，每块卡上的矩阵越小、效率越低，点对点的次数和延迟也越多——这些都不在我们的模型里。结论 #2 是”气泡公式 + 模型之外的代价”一起给出的：模型并行的通信在每个微批、每一层的关键路径上，数据并行的通信每步一次、还能和反向重叠。

7. 小结 · Summary

1. 模型并行把模型的计算切到几块卡上。按层切（流水线）通信最少，可以跨机器；按矩阵切（张量）每层都要 all-reduce，只能在机器内部；两者都让每块卡只存一部分模型。
2. 朴素的流水线利用率只有 1/K。GPipe 把批切成 M 个微批，气泡比例降到 (K − 1)/(M + K − 1)；M 也不能太小，最佳约 √((K − 1)βN/α)。
3. HW5 的 clock_cycles：第 k 波 = {(i, j) : i + j = k}，共 M + K − 1 波。Pipe 一波之内先全部提交再接收；结果与顺序执行数学上相同、浮点上只差舍入。
4. 1F1B 不改变气泡，但把在途激活从 M 降到 K；交错式把气泡缩小 v 倍，代价是 v 倍的传输和更多的在途激活。激活重计算以约 1/3 的额外计算把激活从 O(n) 降到 O(√n)，梯度逐位不变。
5. 张量并行：MLP 先列后行、注意力按头、输出层按词表并融合交叉熵；f / g 两个算子使每层前向、反向各 2 次 all-reduce。序列并行把 LayerNorm 等的激活也切开，通信量不变。
6. 3D 并行：t ≤ 8 在机器内，p 跨机器，t·p 够用就好，其余做数据并行。

8. 常见误区 · Common pitfalls

1. “流水线并行改变了训练结果。” 微批只改变计算的时间安排；梯度是同样的和，只有浮点求和顺序不同（§3.2）。
2. “1F1B 比 GPipe 快。” 气泡完全相同，省的是激活显存（§4.1）；段间传输慢时 1F1B 的关键路径反而更长。
3. “M ≥ 4K 时气泡可以忽略。” 按公式还有 16–18%；GPipe 的说法还依赖重计算填进气泡（§2.3）。
4. “微批越多越好。” 每个微批有固定开销，一步的时间先降后升（§2.4）。
5. “张量并行的 MLP 两个矩阵都按列切。” 第二个必须按行切，才能接住按列分好的中间结果、只用一次 all-reduce；第一个按行切则 GeLU 之前就要通信（§5.1–5.2）。
6. “f 的反向是恒等。” f 前向恒等、反向 all-reduce；g 正好相反（§5.4）。写反了，前向结果对、梯度错——只看前向检查不出来。
7. “行切分层的偏置每块卡都加一次。” 那会加 t 次；偏置要在 all-reduce 之后加一次。
8. “张量并行可以跨机器扩展。” 通信 / 计算之比与批大小无关，跨机器时通信远超计算（§5.6）。
9. “重计算会改变梯度。” 同样的输入、同样的运算，结果逐位相同（§4.3）。
10. “线程里的 torch.no_grad() 自动继承。” 梯度开关是线程局部的（§3.3）。

9. 练习 · Exercises

1. 写出 clock_cycles(2, 3) 的全部输出。一共几波？哪几波里阶段 0 是空闲的？
2. K = 8、M = 24 时，GPipe 的气泡比例是多少？要把气泡降到 10% 以下，M 至少要多少？
3. 一个流水线 K = 4，每个微批每段的前向耗时 t(b) = 0.2 + 0.01·b（ms，b 是行数），反向 = 2 × 前向，批 N = 1024 行。估算 M = 4 和 M = 64 时一步的时间，并用 M* 的公式求最佳 M。
4. 1F1B 下，K = 6 时卡 0、卡 5 最多同时存几个微批的激活？与 M 有关吗？
5. 交错式调度 K = 4、v = 4、M = 16：用 Megatron 的公式估算气泡（相对有效计算）。和 v = 1 比呢？代价是什么？
6. 一串 100 层的网络用每 k 层一个检查点。存的激活个数大约是多少（写成 k 的函数）？k 取多少最省？这时多算了多少次层的前向？
7. 张量并行 t = 4 的 MLP：A 是 1024 × 4096，B 是 4096 × 1024。每块卡上 A、B 的形状是什么？一次前向传多少个数的 all-reduce（输入是 2048 × 1024）？
8. 解释为什么 f 的反向是 all-reduce。如果你把 f 和 g 的反向写反了，哪些梯度会错？
9. 词表 V = 64000、t = 8、一个微批 4 × 4096 个 token：比较 all-gather logits（BF16）与融合交叉熵（3 个 float32 标量）每块卡的通信量。
10. 用 §5.6 的公式说明：为什么把 batch 加大不能改善张量并行的通信 / 计算之比？什么能改善它？
11. 一个 70B 的模型（80 层、h = 8192）要在 256 块 80 GB 的卡（每台 8 块）上训练，训练状态 16 字节 / 参数。只考虑训练状态，t·p 至少多大？你会怎样选 (t, p, d)，为什么？
12. HW5 的 Pipe 在 GPU 上为什么要把 .to(device) 放进提交的函数里，而不是在调用线程里做？

参考答案

1. 4 波：[(0,0)]、[(1,0),(0,1)]、[(1,1),(0,2)]、[(1,2)]。阶段 0 在第 2、3 波空闲（微批只有 2 个）。
2. (8 − 1)/(24 + 7) = 7/31 ≈ 22.6%。7/(M + 7) < 0.1 → M > 63，M ≥ 64。
3. M = 4：每微批 256 行，前向 2.76 ms，一步 = 3 × (4 + 3) × 2.76 = 58.0 ms。M = 64：每微批 16 行，前向 0.36 ms，一步 = 3 × 67 × 0.36 = 72.4 ms。M* = √(3 × 0.01 × 1024 / 0.2) ≈ 12.4（取 M = 12 或 16：M = 16 时 3 × 19 × 0.84 = 47.9 ms）。
4. 卡 0：K − 0 = 6 个；卡 5：1 个。与 M 无关（M ≥ K 时）。
5. (1/v)·(K − 1)/M = (1/4)·(3/16) ≈ 4.7%；v = 1 时 3/16 ≈ 18.8%。代价：跨设备传输约 4 倍，在途激活增加，M 要是 K 的倍数。
6. 约 100/k + k 个；k = 10 时最少（约 20 个）。多算约 100 − 100/k = 90 次层的前向（每段重算 k − 1 层），约多一次完整的前向。
7. A_i：1024 × 1024（按列切 4096 → 1024），B_i：1024 × 1024（按行切 4096 → 1024）。all-reduce 一次，2048 × 1024 = 2,097,152 个数。
8. 每块卡用自己那份权重 A_i 只算出 ∂L/∂X 的一项 ∂L/∂Y_i · A_iᵀ，完整的梯度是 t 项之和。写反了：前向的 g 不做 all-reduce 时连前向都错；若只是反向写反（f 的反向恒等、g 的反向 all-reduce），前向结果正确，但 g 的反向把相同的梯度加了 t 倍，各卡对输入的梯度只是一部分——LayerNorm 和更前面所有层的梯度都会错。
9. all-gather：4 × 4096 × 64000 × 7/8 × 2 = 1.83 GB；融合：16384 × 3 × 4 × 2 × 7/8 = 0.34 MB，相差约 5300 倍。
10. 计算 ∝ b·s·h²/t，通信 ∝ b·s·h，比值 ∝ t/h，与 b、s 无关。能改善它的是更大的 h（更”宽”的模型）、更小的 t、或更快的互连。
11. 70B × 16 = 1120 GB，每卡 80 GB → t·p ≥ 14，取 16（再留余量取 32）。选 t = 8（机器内），p = 2 或 4，d = 256/(t·p) = 16 或 8；激活用重计算 / 序列并行控制。t 不超过 8 是因为跨机器张量并行太慢；p 尽量小，让 d 大、气泡小。
12. 在调用线程里拷贝时，发射线程要等拷贝完成才能发射下一个任务，各阶段的工作被串行化；放进提交的函数里，拷贝在各自阶段的线程（和 CUDA 流）里进行，能和其他阶段的计算重叠。HW5 提示这样做才拿得到 1.10× 的加速。

10. 延伸阅读 · Further reading

• Huang et al. (2019) GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism：微批次、重计算、气泡分析。
• Narayanan et al. (2019) PipeDream: Generalized Pipeline Parallelism for DNN Training（SOSP）：1F1B；以及 PipeDream-Flush / PipeDream-2BW 的后续工作。
• Shoeybi et al. (2019) Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism：张量并行的切法与 f / g 算子。
• Narayanan et al. (2021) Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM（SC）：交错式调度、3D 并行与结论、气泡与 d 的关系、万亿参数的实测。
• Korthikanti et al. (2022) Reducing Activation Recomputation in Large Transformer Models：序列并行、选择性重计算、激活显存公式。
• Chen et al. (2016) Training Deep Nets with Sublinear Memory Cost：√n 检查点。
• Simon Boehm, “Pipeline-Parallelism: Distributed Training via Model Partitioning” 与 shallowspeed 代码（slides 第 6、18–20 页的来源）。
• PyTorch 文档 torch.distributed.pipelining：PipelineStage、ScheduleGPipe / Schedule1F1B / ScheduleInterleaved1F1B。

11. 下一讲预告 · Next lecture

这一讲的模型并行都是”稠密”的：每个 token 都要经过模型的每一个参数。第 17 讲换一个思路——混合专家模型（MoE）：把 FFN 换成很多个”专家”，每个 token 只走其中几个。Mixtral 8×7B 有 47B 参数，每个 token 只激活 13B；DeepSeek-V3 有 671B 参数，每个 token 只激活 37B。参数涨了十几倍，计算量却没怎么涨。代价是一种新的并行——专家并行，以及它带来的 all-to-all 通信和负载均衡问题。我们会从零写一个 top-k 路由的 MoE 层，看容量因子、负载均衡损失和 all-to-all 各自在解决什么。