前置讲 P4 · 多卡与通信 · Multi-GPU Communication

这是一篇前置讲，不对应原课的某一次课。原课第 17 讲”分布式训练”的 slides 从”4000 亿参数光是存下来就要约 800 GB”讲起，依次讲数据并行、张量并行（列切分 / 行切分）、流水线并行（GPipe、1F1B、PipeDream）和 ZeRO；第 16 讲 MoE 讲到”每个专家可以放在不同的 GPU 上”。这些方法的差别，归根结底是用哪种集合通信、通信多少数据、能不能和计算重叠。这篇为没有分布式系统背景的读者把通信这一层讲清楚。 依据说明：本讲讲的是通用的集合通信算法和 α-β 代价模型（参考 Thakur, Rabenseifner & Gropp 2005；Patarasuk & Yuan 2009），以及 PyTorch torch.distributed / DDP 的公开文档。硬件带宽取自 NVIDIA H100 / NVLink 4 与 InfiniBand NDR 的公开规格（标称值）。代码里的”多块 GPU”是用列表模拟的，通信算法是真实执行的；05 号脚本另外在本机真正启动 4 个进程，用 gloo 后端跑一次真实的分布式训练。 配套代码 Code：code/prereq4/（6 个脚本 + 1 个模块；不需要 GPU；每个脚本 1–5 秒） 预计阅读时间：3 小时（含运行代码） 建议阅读时机：第 17 讲之前；第 16 讲讲专家并行时也会用到。建议先读前置讲 P1 §5（显存核算）和 P3 §6（训练的 6ND）。

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

0.1 为什么要懂这些

前置讲 P1 算过：用 Adam 训练一个 8B 模型，光是参数、梯度和优化器状态就要约 120 GiB，一块 80 GB 的 H100 放不下。P3 算过：在 15T token 上训练它要约 7×10²³ 次浮点运算，一块 H100 要跑五十多年。所以训练一定要用很多块 GPU，而一旦有了很多块 GPU，就出现了一个新问题：它们之间怎样交换数据？要花多少时间？

这篇回答这个问题。学完之后，你应该能看懂下面这些说法背后的道理：

• “数据并行的通信量只和参数量有关，和批大小无关”（§5）
• “环形 all-reduce 让每块卡的通信量几乎与卡数无关”（§3）
• “PyTorch DDP 默认把梯度分成 25 MB 的桶”（§4）
• “张量并行要放在节点内，数据并行和流水线并行可以跨节点”（§7）
• “FSDP 用 50% 的额外通信换来显存降到 1/N”（§7）
• “MoE 的 all-to-all 是瓶颈，所以要做负载均衡”（§7）

0.2 学习目标 Learning goals

1. 知道多卡训练的基本术语：进程、rank、world size、节点；NVLink、InfiniBand、PCIe 的带宽数量级；NCCL 和 gloo 是什么。
2. 说出六种集合通信（broadcast、reduce、all-reduce、reduce-scatter、all-gather、all-to-all）前后每块卡有什么，以及 all-reduce = reduce-scatter + all-gather。
3. 能一步一步讲清环形 all-reduce，推导每块卡的通信量 $2(N-1)/N \cdot S$，并解释它为什么可扩展。
4. 会用 α-β 模型估计通信耗时，知道小消息被延迟主导、大消息被带宽主导；理解梯度分桶和分层 all-reduce 的动机。
5. 理解数据并行为什么在数学上等价于大批量训练，和梯度累积的关系，以及各卡样本数不均时会出什么错。
6. 会写一个最小的 torch.distributed 程序，知道 DDP 在构造和反向时各做了什么。
7. 能给一个模型的训练算一笔”通信账”：数据并行、FSDP、张量并行、MoE 各要通信多少、和计算时间比如何。

0.3 本讲地图 Map

§1  为什么一块 GPU 不够；术语与硬件                          
§2  六种集合通信                                              ← 代码 01
§3  环形 all-reduce：逐步追踪、通信量、数值                   ← 代码 02
§4  α-β 模型：消息大小、环形 vs 树形、分桶、分层              ← 代码 03
§5  数据并行：和大批量等价、梯度累积、不均匀的分片            ← 代码 04
§6  真正的分布式程序：torch.distributed 与 DDP               ← 代码 05
§7  通信账：数据并行、FSDP、张量并行、MoE                     ← 代码 06
§8  小结   §9  常见误区   §10 练习   §11 延伸阅读   §12 下一步

1. 为什么一块 GPU 不够 · Why one GPU is not enough

1.1 两堵墙

1. 显存墙：8B 模型的训练状态约 120 GiB（P1 §5.4），70B 约 1 TiB，还不算激活值。一块 H100 只有 80 GB。
2. 算力墙：8B 模型在 15T token 上训练要约 51 万 H100 小时（P3 §6.4，MFU 40%）。一块卡做不完，一千块卡也要三周。

两堵墙都只有一个办法：把工作分给很多块 GPU。分法有好几种（第 17 讲详细讲）：

无论哪种切法，各块卡之间都必须交换数据：梯度、参数、激活值、token。通信是多卡训练的核心开销。

1.2 术语

1.3 硬件：卡和卡之间有多快

两个数量级的差距要记住：节点内 NVLink 比节点间 InfiniBand 快约 9 倍；显存又比 NVLink 快约 7 倍。这决定了”哪种并行放在哪一层”的设计（§7）。

2. 六种集合通信 · The six collectives

01 用 4 块”卡”演示，rank $r$ 的数据是 $[10r+1, 10r+2, 10r+3, 10r+4]$，一眼能看出每个数来自哪块卡：

（all-to-all 相当于把一个”rank × 块”的矩阵转置。）

最重要的一个关系：

reduce-scatter 之后每人有总和的一块，all-gather 再把这些块传给所有人——合起来就是每人都有完整的总和。这个分解是理解 ZeRO / FSDP 的钥匙：它们把数据并行的一次 all-reduce 拆成两半，在不同的时机分别做（§7.2）。

每块卡收发多少（01 第 2 部分，S = 每块卡的数据量）：

谁用哪种通信：

3. 环形 all-reduce · Ring all-reduce

3.1 朴素做法的问题

最直接的 all-reduce：所有人把数据发给 rank 0，rank 0 求和，再发回给所有人。rank 0 要收 $(N-1)S$、发 $(N-1)S$，其他人都在等它——根节点是瓶颈，耗时随卡数线性增长。

3.2 算法

环形 all-reduce（ring all-reduce）把 N 块卡排成一个环，每块卡只和左右邻居通信。每块卡的数据切成 N 块。分两个阶段：

1. reduce-scatter（N−1 步）：第 $s$ 步，rank $r$ 把它手里第 $(r - s) \bmod N$ 块的”部分和”发给右边的邻居，邻居把它加到自己的同一块上。一块数据沿着环走一圈，经过每块卡时都加进了那块卡的贡献；N−1 步后，每块卡各有一块完整的总和。
2. all-gather（N−1 步）：每块卡把刚拿到的完整块传给右边的邻居，N−1 步后所有人都有了全部 N 块。

02 第 1 部分用一个巧妙的初值追踪这个过程：rank $r$ 的 4 块都初始化为 $10^r$（rank 0 = 1，rank 1 = 10，rank 2 = 100，rank 3 = 1000）。一个块的值哪几位是 1，就说明它已经包含了哪几块卡的贡献，加满是 1111：

第 1 步（reduce-scatter）后：rank 0:    1     1     1  1001     ← rank 3 把第 3 块传给了 rank 0
第 3 步（reduce-scatter）后：rank 0:    1  1111  1101  1001     ← rank 0 拥有完整的第 1 块
                             rank 1:   11    10  1111  1011
                             rank 2:  111   110   100  1111
                             rank 3: 1111  1110  1100  1000
第 6 步（all-gather）后：    每块卡: 1111  1111  1111  1111

每一步，所有链路同时在工作，每块卡同时发一块、收一块，没有谁比别人更忙。

3.3 通信量：几乎与卡数无关

每一步每块卡发 $S/N$，共 $2(N-1)$ 步，所以每块卡收发

02 第 2 部分用模拟器逐步记账验证：

环形算法里每块卡的通信量最多不超过 2S，和卡数几乎无关。可以证明这是 all-reduce 在带宽意义上的下界（Patarasuk & Yuan, 2009）：每块卡至少要把自己的 $(N-1)/N$ 份数据送出去参与求和，再把 $(N-1)/N$ 份结果收回来。

在 α-β 模型（§4）下，1 GiB 数据在 NVLink 上做环形 all-reduce：2 块卡 2.40 ms，8 块卡 4.25 ms，256 块卡 7.30 ms——基本是 $2S/\beta \approx 4.77$ ms 附近；朴素方案 256 块卡要 1.2 秒。代价是步数：环形要 $2(N-1)$ 步，每步有一个固定延迟 $\alpha$，卡数很多时（256 块卡时延迟项 2.55 ms）开始显现，这时会改用树形或分层算法（§4）。

3.4 数值：所有卡的结果逐位相同

02 第 4 部分：8 块卡、每块 4096 个 FP32 随机数。环形算法里每一块的求和顺序取决于它从环上哪块卡出发，和”按 rank 0、1、2…逐个相加”不同，所以有 52% 的元素在最后几位上有差别（P1 §2.6：浮点加法不满足结合律）。但 all-gather 阶段只是复制，所以所有卡拿到的结果逐位相同。这一点非常重要：如果各卡的梯度有哪怕最微小的差别，各卡上的模型参数就会慢慢”分叉”，数据并行就不成立了。

另一方面，换卡数、换通信算法都会改变求和顺序，所以”同样的代码、不同的集群规模”训练出来的模型不能逐位复现。

4. α-β 模型 · The alpha-beta cost model

4.1 公式

一次传输 $n$ 字节的耗时：

$\alpha$ 是固定延迟（启动一次传输、经过软件栈和网络的开销，微秒量级），$\beta$ 是带宽。环形 all-reduce 有 $2(N-1)$ 步、每步传 $S/N$：

03 第 1 部分验证：模拟器把每一步”最忙的那块卡收发多少字节”换算成时间再加起来，和公式完全一致（4 卡 16 MiB 时 85.9 µs）。下面用 $\alpha = 5$ µs（粗略的量级）、NVLink $\beta = 450$ GB/s。

4.2 小消息看延迟，大消息看带宽

树形 all-reduce（沿二叉树先 reduce 再 broadcast）只要 $2\lceil\log_2 N\rceil$ 步，但每步传整份数据：$T_{\text{tree}} \approx 2\lceil\log_2 N\rceil(\alpha + S/\beta)$。03 第 2 部分（8 块卡）：

• 小消息：耗时几乎全是延迟，步数少的树形更快。
• 大消息：带宽项占主导，每步只传 $1/N$ 的环形快得多。
• 有效带宽只有在大消息时才接近链路带宽。所以通信要”攒大”了再发；NCCL 会根据消息大小、卡数和拓扑自动选择环形或树形算法（真实的树形算法会把数据切块流水线化，比这里的简单版本快）。

4.3 梯度分桶

数据并行要 all-reduce 全部梯度。03 第 3 部分：1 GiB 梯度、8 块卡，把梯度切成若干”桶”分别 all-reduce：

桶太小，延迟吃掉一切。那为什么不干脆用一个最大的桶？因为分桶的真正目的是让通信和计算重叠：反向传播从最后一层往前算，最后几层的梯度一算好，就可以先把这一桶 all-reduce 出去，同时继续算前面层的梯度。桶太大，要等很久才能开始通信。PyTorch DDP 默认 25 MB（bucket_cap_mb）是折中：在这个模型里它让总通信时间比一个大桶多约 70%，但这些通信大部分藏在反向计算后面，不在关键路径上。最优的桶大小取决于延迟、带宽和计算速度。

4.4 跨节点：分层 all-reduce

一个大集群由很多节点组成，节点内 NVLink（450 GB/s），节点间 InfiniBand（每块卡约 50 GB/s）。如果把所有卡排成一个平坦的大环，每一步都要等最慢的那段链路，整个 all-reduce 被拖慢到 InfiniBand 的速度。

分层 all-reduce：

1. 节点内 reduce-scatter（NVLink）：每块卡得到节点内部分和的 1/8；
2. 跨节点 all-reduce（InfiniBand）：各节点的”同号卡”（都负责第 $k$ 块的那些卡）组成一个环，只对这 1/8 的数据做 all-reduce；
3. 节点内 all-gather（NVLink）。

慢链路上传的数据减少到 1/8。03 第 4 部分（1 GiB）：

这个思路贯穿大模型训练的并行设计：把通信量大、又不能重叠的并行方式放在快链路上（节点内），把通信量小或能重叠的放在慢链路上（节点之间）（§7.3）。

5. 数据并行 · Data parallelism

5.1 算法

数据并行（data parallelism）：每块卡有一份完整的模型；每一步，全局的一批样本平均分给 N 块卡；每块卡在自己那份上做前向、反向，得到本地的平均梯度；all-reduce 求和再除以 N；每块卡用同样的平均梯度更新自己的模型副本。

5.2 它等价于大批量训练

设全局批大小 $B$，每块卡分到 $B/N$ 个样本。本地平均梯度 $g_r = \frac{N}{B}\sum_{i \in \text{第 } r \text{ 份}} \nabla \ell_i$，于是

正是全局批的平均梯度。04 用 softmax 回归、全局批 256、50 步验证：

差异只是求和顺序带来的舍入（§3.4）。各卡参数逐位相同：all-reduce 保证每块卡拿到完全一样的梯度，同样的初始化 + 同样的更新 = 同样的参数。所以 DDP 只在开始时广播一次参数，之后不需要再同步参数。

5.3 通信量只和参数量有关

每一步每块卡收发 $2(N-1)/N \times$ 参数字节（04 实测与公式一致）——和批大小无关。批越大，每步的计算越多，而通信不变，通信的相对开销就越小。所以数据并行适合”每块卡上计算量大、参数相对少”的情况；每卡的批太小时，通信会成为瓶颈（§7.1）。

5.4 梯度累积：在时间上做数据并行

梯度累积（gradient accumulation）：一块卡把一个大批切成 $k$ 个小批，逐个前向、反向，把梯度加起来，最后才更新一次。04 第 2 部分：4 个小批的梯度累积和 4 块卡的数据并行、和单卡一次算完的结果都只差 2.2e-16。数据并行在空间上分批，梯度累积在时间上分批，算的是同一个东西。显存放不下大批量时，就用梯度累积”假装”有更多的卡——P1 和第 1 讲说的”PyTorch 的梯度会累加”就是为它准备的。

5.5 各卡样本数不均时会出错

04 第 3 部分：每步 256 个样本按 [100, 60, 60, 36] 分给 4 块卡，各卡求自己的平均梯度后再简单平均——结果和单卡差了 2.8e-2。“平均的平均”不再等于全局平均：样本少的卡被赋予了过大的权重（36 个样本的卡和 100 个样本的卡权重一样）。所以 PyTorch 的 DistributedSampler 会让每块卡分到同样多的样本（必要时补齐或丢弃最后一批）。

6. 真正的分布式程序 · torch.distributed and DDP

6.1 最小的分布式程序

05 在本机真正启动 4 个进程，每个进程扮演一块”卡”，通过本机回环网络通信。结构是：

import torch.distributed as dist

def worker(rank, world_size, port):
    os.environ["MASTER_ADDR"] = "127.0.0.1"          # 所有进程在这里"会合"
    os.environ["MASTER_PORT"] = str(port)
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    x = torch.arange(1, 5.) + 10 * rank
    dist.all_reduce(x)                                # 原地求和：每个进程都得到总和
    dist.destroy_process_group()

torch.multiprocessing.spawn(worker, args=(4, port), nprocs=4)

真实的多 GPU 训练把 "gloo" 换成 "nccl"，并用 torchrun --nproc_per_node=8 train.py 启动（它负责设置 RANK、LOCAL_RANK、WORLD_SIZE 等环境变量）；编程接口完全相同。

05 第 1 部分用 all_reduce、broadcast、all_gather、reduce_scatter、all_to_all_single 重做了 §2 的例子，结果和用列表模拟的完全一致（总和 [64, 68, 72, 76]）。一个小差别：PyTorch 的 reduce_scatter 让 rank $r$ 拿总和的第 $r$ 块，comm_lab 的环形实现让 rank $r$ 拿第 $r+1$ 块——只是”谁负责哪一块”的编号不同。

一条铁律：每个集合通信都必须由所有 rank 以同样的顺序调用。只要有一个 rank 少调了一次（例如写了 if rank == 0: dist.all_reduce(x)），其他 rank 就会永远等下去——这是分布式程序最常见的 bug：死锁（hang）。

6.2 DDP 做了什么

05 第 2 部分：4 个进程，故意让每个进程的初始参数都不同，再用 DistributedDataParallel 包装：

• 构造时：DDP 把 rank 0 的参数广播给所有人。包装之前 4 个 rank 的权重各不相同，包装之后完全一样。
• 反向时：DDP 在每个参数上注册了钩子；梯度算好一个桶，就立刻 all-reduce 这个桶并除以 world size（§4.3 的分桶和重叠）。
• 结果：训练 20 步后，4 个 rank 的权重逐位相同；和单进程用 256 个样本的大批量训练相比，最大差异 1.2e-7（FP32 舍入）——和 §5 的模拟结论一致。

6.3 真实的延迟和带宽

05 第 3 部分在本机回环网络上测 all-reduce（⏱，一次运行的结果）：1 KiB 时一次要约 880 µs，几乎全是固定延迟（进程间要经过 TCP 协议栈）；64 MiB 时总线带宽约 3.3 GB/s。形状和 α-β 模型一致：小消息看延迟，大消息看带宽。但 gloo 走 CPU 和网络协议栈，延迟和带宽都比 NVLink 上的 NCCL 差几十到上百倍——只用它来理解接口和算法，不要用它的数字推断 GPU 集群的性能。

7. 通信账：训练一个 8B 模型 · A communication budget

06 以 Llama-3 8B（BF16 梯度 15 GiB）在 H100 集群上训练为例，用 α-β 模型和 roofline（MFU 40%）算账。全部是理想估计，只看量级和趋势。

7.1 数据并行

• 每卡每步 32K token 时，梯度通信只占计算的几个百分点，而且能和反向重叠，几乎是”免费”的。
• 跨节点后通信时间约翻倍，但从 64 卡到 1024 卡几乎不再增长——分层环形算法让每块卡的通信量和卡数基本无关。
• 每卡的批变小（4K token），计算按比例缩短、通信不变，占比升到 12%–27%：每卡的批不能太小。

7.2 FSDP：用通信换显存

数据并行要求每块卡放下完整的训练状态：8B 模型 120 GiB，H100 放不下。FSDP（Fully Sharded Data Parallel，ZeRO-3 的 PyTorch 实现）把参数、梯度、优化器状态都切成 N 份，每块卡只存自己那份：

代价是通信：每一层计算前要 all-gather 这一层的完整参数（用完就丢），前向一次、反向再一次；反向后 reduce-scatter 梯度，每块卡只拿到自己那 1/N 份的梯度、只更新自己那 1/N 的参数和优化器状态。合计 $3 \times (N-1)/N$ 份参数量，而 DDP 的一次 all-reduce 是 $2 \times (N-1)/N$ 份。FSDP 多了 50% 的通信，换来显存降到 1/N——这就是 §2 说的”把 all-reduce 拆成两半”。第 17 讲会讲 ZeRO 的三个阶段（只切优化器状态、再切梯度、再切参数）。

7.3 张量并行：只放在节点内

张量并行（Megatron 式）把每一层的矩阵切成 $t$ 份。每一层前向要 all-reduce 两次激活（注意力输出、MLP 输出），反向再两次。06 第 3 部分：TP = 8，8 块卡一起处理 32K token，每次 all-reduce 的激活 256 MiB，每步 128 次：

两个关键差别：(1) 张量并行的通信量和 token 数成正比（传的是激活），而不是和参数量；(2) 这些 all-reduce 夹在每一层的计算中间，后面的计算要等它完成，很难整体重叠。所以张量并行几乎总是限制在一个节点（8 块卡）之内，跨节点用数据并行和流水线并行——这就是第 17 讲的”3D 并行”。

7.4 MoE 的 all-to-all

专家并行把 MoE 的不同专家放在不同的卡上。每一层 MoE 都要做两次 all-to-all：把每个 token 送到它的专家所在的卡，算完再取回来；反向时再做一遍。06 第 4 部分（专家分在 8 块卡上、每卡每步 8192 个 token、top-1）：每层每卡收发 112 MiB，节点内每层约 0.27 ms、跨节点约 2.4 ms，32 层前向 + 反向在跨节点时约 226 ms。

all-to-all 同样夹在计算中间。更麻烦的是负载不均衡：如果路由器把大量 token 送给某一个专家，那块卡要收的数据特别多，all-to-all 的耗时由最忙的那块卡决定——和 §3.1 朴素方案里根节点的瓶颈是同一个道理。这是第 16 讲引入负载均衡损失的系统层面的动机。

8. 小结 · Summary

1. 一块 GPU 放不下（8B 模型训练状态 120 GiB）、也算不完（51 万 GPU 小时）大模型的训练，必须用多卡；多卡的核心开销是通信。节点内 NVLink 约 450 GB/s，节点间 InfiniBand 约 50 GB/s。
2. 六种集合通信：broadcast、reduce、all-reduce、reduce-scatter、all-gather、all-to-all。all-reduce = reduce-scatter + all-gather。
3. 环形 all-reduce：2(N−1) 步，每块卡收发 $2(N-1)/N \cdot S < 2S$，与卡数几乎无关（朴素方案的根节点要 $(N-1)S$）；所有卡的结果逐位相同，但求和顺序与串行不同。
4. α-β 模型 $T = \alpha + n/\beta$：小消息被延迟主导（树形更快），大消息被带宽主导（环形更快）。梯度分桶是在延迟和”通信与计算重叠”之间折中（DDP 默认 25 MB）；跨节点用分层 all-reduce，让慢链路只传 1/8 的数据（64 卡时快 4.8 倍）。
5. 数据并行在数学上等价于大批量训练（FP64 差 2.2e-16），各卡参数逐位相同；通信量只和参数量有关、和批大小无关；梯度累积是”时间上的数据并行”；各卡样本数必须相同。
6. torch.distributed：init_process_group、all_reduce 等接口，所有 rank 必须以同样顺序调用集合通信；DDP 构造时广播参数、反向时按桶 all-reduce 梯度并与计算重叠。真实的 4 进程 DDP 与单进程大批量差 1.2e-7。
7. 通信账（Llama-3 8B）：数据并行的梯度通信在每卡批足够大时只占计算的几个百分点；FSDP 用 50% 的额外通信把显存降到 1/N；张量并行的通信与 token 数成正比、夹在计算中间，跨节点会比计算还慢，所以留在节点内；MoE 的 all-to-all 受最忙的卡拖累，需要负载均衡。

9. 常见误区 · Common pitfalls

1. “卡数翻倍，每块卡的通信量也翻倍。” 环形 all-reduce 每块卡的通信量最多 2S，和卡数几乎无关。
2. “all-reduce 就是先 reduce 再 broadcast。” 语义上没错，但那是最慢的实现；高效的实现是 reduce-scatter + all-gather（环形）或树形算法。
3. “数据并行训练出来的模型和单卡不一样。” 数学上等价（各卡样本数相同时），只差舍入误差；各卡之间参数逐位相同。
4. “批大小越大，数据并行的通信越多。” 通信量只和参数量有关；批越大，通信的相对开销越小。
5. “把梯度攒成一个大桶通信最快。” 单看通信是这样，但会失去和反向计算重叠的机会。
6. “各卡分到的样本数差一点没关系。” 简单平均各卡的平均梯度会给小分片过大的权重；04 里差了 2.8e-2。
7. “只有 rank 0 需要调用 all_reduce。” 所有 rank 都必须调用，而且顺序要一致，否则死锁。
8. “张量并行和数据并行一样，跨节点也没问题。” 张量并行的通信与 token 数成正比、夹在每一层中间，跨节点会比计算还慢。
9. “FSDP 只是更省显存的 DDP，没有代价。” 通信量多 50%，而且参数的 all-gather 要精心安排预取才能和计算重叠。
10. “gloo 测出来的通信速度能代表 GPU 集群。” gloo 走 CPU 和 TCP，比 NVLink 上的 NCCL 慢几十到上百倍。

10. 练习 · Exercises

练习 1（通信量）。 16 块卡，每块卡有 4 GiB 的梯度。(a) 环形 all-reduce 中每块卡发出多少数据？(b) 朴素方案中根节点要收多少？(c) 在 450 GB/s 的链路上，忽略延迟，两者各要多久？

(a) $2 \times 15/16 \times 4 = 7.5$ GiB。(b) $15 \times 4 = 60$ GiB（收）+ 60 GiB（发）。(c) 环形约 $7.5 \times 2^{30} / 450 \times 10^9 \approx 17.9$ ms；朴素方案根节点收 60 GiB 约 143 ms、再发 60 GiB 约 143 ms，合计约 286 ms，慢 16 倍。

练习 2（环形算法的步数）。 (a) 为什么环形 reduce-scatter 需要恰好 N−1 步？(b) 如果每块卡的数据不切成 N 块，而是整份沿环传递（“一圈一圈地加”），通信量是多少？

(a) 每一块数据要经过另外 N−1 块卡才能收集到所有贡献；每一步它只走一格，所以要 N−1 步。所有块同时在环上流动，互不等待。

(b) 每一步每块卡传整份 S，要 N−1 步才能让总和走完一圈，再 N−1 步把结果传回所有人：每块卡收发约 $2(N-1)S$，是切块版本的 N 倍。切块的意义就是让每一步只传 $1/N$，所有链路同时满载。

练习 3（环形 vs 树形的转折点）。 用 §4.2 的两个公式，求 8 块卡、$\alpha = 5$ µs、$\beta = 450$ GB/s 时，消息多大时环形和树形一样快。和 03 第 2 部分的表一致吗？

$14(\alpha + S/(8\beta)) = 6(\alpha + S/\beta)$，即 $8\alpha = (6 - 1.75)S/\beta$，$S = 8\alpha\beta / 4.25 \approx 4.2 \times 10^6$ 字节 ≈ 4.0 MiB。表里 1 MiB 时树形更快、16 MiB 时环形更快，转折点落在两者之间，一致。用 comm_lab.alpha_beta_ring_allreduce 和 alpha_beta_tree_allreduce 代入 4.0 MiB 可以验证两者相等（本机算出都是 86.5 µs）。

练习 4（大模型的数据并行）。 Llama-3 70B 的 BF16 梯度约 141 GB。在 8 个节点 × 8 块 H100 上做分层 all-reduce 要多久？如果每块卡每步处理 32K token，计算时间约多少（MFU 40%）？通信占比多少？（提示：可以直接调用 comm_lab.hierarchical_allreduce_time。）

hierarchical_allreduce_time(70.55e9 * 2, 8, 8, 5e-6, 450e9, 50e9) 本机算出约 1.17 s（节点内两段各约 0.27 s、跨节点约 0.62 s）。计算时间 $6 \times 70.55 \times 10^9 \times 32768 / (0.4 \times 989 \times 10^{12}) \approx 35$ s，通信占比约 3%。所以 70B 模型的数据并行在通信上没问题——问题在于每块卡放不下 1 TiB 的训练状态，必须配合 FSDP / ZeRO 或模型并行。

练习 5（梯度累积的代价）。 一块卡想模拟 8 块卡的数据并行（全局批 2048，每块卡 256），用梯度累积要怎么做？显存和时间各有什么变化？

把 2048 个样本分成 8 个 256 的小批，逐个前向、反向，梯度累加（不清零），8 个小批之后除以 8（或者每个小批的损失先除以 8）再更新一次。显存只需放下 256 个样本的激活（和真的 8 卡时每块卡一样），参数和优化器状态只有一份；时间是 8 卡的约 8 倍（没有并行）。它是”用时间换空间”，结果在数学上和 8 卡数据并行相同。

练习 6（不均匀分片的正确做法）。 如果各卡的样本数确实不同（$n_1, \dots, n_N$），怎样用 all-reduce 得到正确的全局平均梯度？

每块卡算梯度之和（而不是平均）$G_r = \sum_{i \in r} \nabla\ell_i$，同时把样本数 $n_r$ 也放进一个张量；all-reduce 两者，得到 $\sum_r G_r$ 和 $\sum_r n_r$，相除就是全局平均。等价地，每块卡先算平均梯度再乘以 $n_r \cdot N / \sum_r n_r$ 的权重。处理变长序列时（每块卡的 token 数不同），按 token 数加权也是同样的道理。

练习 7（哪种通信随什么增长）。 填表：数据并行、FSDP、张量并行、专家并行的通信量分别和什么成正比（参数量 / token 数）？它们能不能和计算重叠？

规律：和参数量成正比、能重叠的放在慢链路上（跨节点），和 token 数成正比、不能重叠的放在快链路上（节点内）。

练习 8（FSDP 的 3 倍）。 推导 FSDP 每一步每块卡的通信量 $3(N-1)/N \cdot P$（$P$ = 参数字节数）。如果显存足够、反向时不重新 all-gather 参数（前向 all-gather 后一直留着），通信量是多少？显存呢？

前向：每一层 all-gather 参数，$(N-1)/N \cdot P$；反向：再 all-gather 一次（前向后已丢弃），$(N-1)/N \cdot P$；梯度 reduce-scatter，$(N-1)/N \cdot P$。合计 $3(N-1)/N \cdot P$。如果前向后保留完整参数，就省掉反向那次 all-gather，通信量 $2(N-1)/N \cdot P$（和 DDP 一样），但每块卡要放下完整的 BF16 参数（这就是 ZeRO-2 的做法：只切梯度和优化器状态）。显存与通信的取舍是 ZeRO 各阶段的核心。

练习 9（死锁）。 下面的代码会发生什么？怎样改？

loss = compute_loss()
if dist.get_rank() == 0:
    dist.all_reduce(loss)
    print("平均损失", loss.item() / dist.get_world_size())

rank 0 调用了 all_reduce，在等其他 rank 加入；其他 rank 没有调用，直接往下走，可能在下一个集合通信（例如 DDP 反向时的梯度 all-reduce）里等待——两边等的不是同一个操作，结果是永久挂起，或者更糟：错配成两个不同的操作、得到错误的结果。改法：所有 rank 都调用 dist.all_reduce(loss)，只把 print 放在 if rank == 0 里。

练习 10（负载不均衡）。 MoE 有 8 个专家分在 8 块卡上。如果路由器把 25% 的 token 送给专家 0（平均应该是 12.5%），all-to-all 的耗时大约是均衡时的几倍？

all-to-all 的耗时由收发最多的那块卡决定。专家 0 所在的卡要收的数据是均衡时的 2 倍（25% / 12.5%），所以耗时约为 2 倍；而且专家 0 的计算量也翻倍，其他卡在等它。这就是为什么 MoE 训练要加负载均衡损失、设”容量因子”上限（超过容量的 token 被丢弃或改走残差路径），第 16 讲详细讨论。

11. 延伸阅读 · Further reading

• Thakur, Rabenseifner & Gropp (2005), Optimization of Collective Communication Operations in MPICH：环形、树形、递归减半等集合通信算法及其 α-β 代价。
• Patarasuk & Yuan (2009), Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations：环形 all-reduce 的带宽最优性。
• Sergeev & Del Balso (2018), Horovod: fast and easy distributed deep learning in TensorFlow：把环形 all-reduce 带进深度学习的工作之一。
• Li et al. (2020), PyTorch Distributed: Experiences on Accelerating Data Parallel Training, VLDB：DDP 的设计（分桶、通信与计算重叠）。
• Rajbhandari et al. (2020), ZeRO: Memory Optimizations Toward Training Trillion Parameter Models；Zhao et al. (2023), PyTorch FSDP。
• Shoeybi et al. (2019), Megatron-LM：张量并行的列切分与行切分（第 17 讲的阅读材料）。
• NVIDIA NCCL 文档与 nccl-tests：真实硬件上集合通信性能的测量方法（总线带宽 busbw 的定义）。
• 本仓库的 CS336 第 7 讲（并行 I） 和 CS349D 前置讲 P5（分布式通信）：用 gloo 多进程实现 DDP / ZeRO，以及推理系统中的通信。

12. 下一步 · What’s next

四篇前置讲到这里就全部完成了：P1 讲数字和内存，P2 讲概率工具，P3 讲单块 GPU 的性能，P4 讲多块 GPU 之间的通信。

• 如果你是按顺序学习，下一步是第 2 讲：Transformer 语言模型（撰写中）；第 1 讲见 lecture01_overview_rnn_autodiff.md。
• P2 的内容从第 6 讲（GAN）开始大量使用；P3、P4 的内容在第 16 讲（MoE）、第 17 讲（分布式训练）、第 18 讲（FlashAttention 与高效解码）、第 19 讲（长上下文）中反复出现，读到那几讲时可以回来查。