第 7 讲：并行 I —— 集合通信、数据并行与 ZeRO

对应：Stanford CS336 Spring 2026 Lecture 7（Parallelism，Percy Liang 的可执行讲义：通信原语、硬件、torch.distributed、数据/张量/流水线并行的最简实现）和 Lecture 8（Tatsunori Hashimoto 的幻灯片）中关于数据并行与 ZeRO 的部分。Spring 2025 的第 7、8 讲内容相近，有公开录像。 配套代码：code/lecture07/ 预计阅读时间：3–4 小时 硬件说明：不需要 GPU。配套代码在一台电脑上启动多个进程，每个进程扮演一块 GPU，用 PyTorch 的 gloo 后端通信。编程接口与真实多 GPU 训练完全相同。

0. 本讲导航

0.1 学完本讲，你应该能够

1. 用显存、算力、通信三笔账说明：为什么训练大模型必须用很多块 GPU。
2. 说出 GPU 之间的连接层次（NVLink、InfiniBand 等）和各级带宽的数量级，以及它对设计并行策略的影响。
3. 说出八种集合通信原语（broadcast、scatter、gather、reduce、all-gather、reduce-scatter、all-reduce、all-to-all）的含义和用途，并解释为什么 all-reduce = reduce-scatter + all-gather。
4. 手推环形 all-reduce，算出它的通信量 $2(p-1)/p \cdot n$，并用”延迟 + 带宽”模型比较不同算法。
5. 用 torch.distributed 写出 SPMD 风格的分布式程序。
6. 从零实现数据并行（DDP），证明它等价于大批次训练，并实现”通信与计算重叠”和”梯度分桶”。
7. 算出混合精度训练中每个参数占 16 字节，并用 ZeRO-1/2/3 的公式估算每块卡的显存。
8. 实现 ZeRO-1（优化器状态分片）和 ZeRO-3 / FSDP（参数也分片），说清楚为什么 ZeRO-1 是”免费的”、ZeRO-3 的通信量为什么是 $3\Psi$。

0.2 先修知识

0.3 本讲地图

第1节 为什么需要多块 GPU：显存、算力、通信三笔账
第2节 硬件：GPU 之间是怎么连起来的
第3节 集合通信原语
第4节 通信的代价：环形 all-reduce、α-β 模型、带宽测量
第5节 用 torch.distributed 编程（以及本讲代码怎么在一台电脑上模拟多卡）
第6节 数据并行（DDP）：等价性、实现、重叠与分桶、极限
第7节 显存：每个参数 16 字节，而且存了 N 份
第8节 ZeRO：ZeRO-1、ZeRO-2、ZeRO-3（FSDP）

第 5、6 讲的主题是”在一块 GPU 内部，让计算不被数据搬运拖慢”；这一讲和下一讲的主题是”在很多块 GPU 之间，让计算不被通信拖慢”。你会发现两者的思想惊人地相似。

1. 为什么需要多块 GPU？

1.1 两个直接的原因

CS336 课上给出的理由很朴素：

1. 放不下：参数、梯度、优化器状态、激活值加起来，一块 GPU 的显存装不下。
2. 太慢了：一块 GPU 的算力有限，想训练得更快，就要更多的算力。

具体算一下（01_why_parallelism.py）。

显存账。第 2 讲算过，用 AdamW 做混合精度训练，每个参数要占 16 字节：

这还没算激活值。结论很直接：模型状态必须切开，分到很多块 GPU 上。

算力账。训练计算量约为 $6ND$（第 2 讲）。假设 H100 的 BF16 峰值 989 TFLOP/s，MFU 为 40%：

Llama 3 405B 的 $3.8\times10^{25}$ 与 Meta 论文报告的数字一致；Meta 实际用了最多 16384 块 H100。一块卡要算三千年的事，只能分给上万块卡一起做。

注意：上表”卡数翻倍、时间减半”的假设是理想情况。现实中，卡越多，通信越多，每块卡的利用率往往越低。整个第 7、8 讲的目标，就是让这个假设尽量成立：显存随卡数线性减少，算力随卡数线性增加（CS336 课件的说法）。

1.2 统一的主题：计算离数据很远

回忆第 5 讲：GPU 内部，计算单元旁边的寄存器、共享内存很快但很小，HBM 很大但慢得多。我们用算子融合和分块减少对 HBM 的读写。

到了多卡，存储层次又往外延伸了好几级（CS336 课上的”广义层次结构”）：

单卡：L1 缓存 / 共享内存（最快）
单卡：HBM
单机多卡：NVLink / NVSwitch
多机多卡：InfiniBand / 以太网（最慢）

统一的主题：计算单元离输入输出数据很远，要精心编排计算，避免被数据搬运卡住。

• 上周（单卡）：用融合 / 分块减少显存访问；
• 这周（多卡）：用**复制（replication）/ 切分（sharding）**减少设备之间的通信。

后面你会看到一个反复出现的权衡：一样东西（参数、激活值），要么自己存一份（占显存），要么用的时候找别人要（花通信），要么重新算一遍（花计算）。CS336 的总结：可以重计算，可以存在显存里，也可以存在别的 GPU 的显存里、需要时通信。

1.3 并行可以沿哪些”维度”切？

训练的计算可以沿四个方向切开，分给不同的 GPU：

这一讲专注于第一种，以及让它省显存的 ZeRO 系列方法。

2. 硬件：GPU 之间是怎么连起来的

2.1 家用 vs 数据中心

家用电脑：同一台机器里的几块显卡通过 PCIe 总线相连，机器之间通过普通以太网相连。PCIe 5.0 x16 单向约 64 GB/s；千兆以太网只有约 0.125 GB/s。

数据中心（CS336 课上描述的典型配置）：

• 一台机器（节点）8 块 GPU，通过 NVLink 连到 NVSwitch，任意两块 GPU 之间都能全速通信。H100 的 NVLink 总带宽约 900 GB/s（双向之和），B200 约 1.8 TB/s；
• 节点之间通过 InfiniBand 相连，通常每块 GPU 配一块 400 Gb/s 的网卡，约 50 GB/s；
• 很多节点组成一个”pod”，很多 pod 组成一个集群 / 数据中心，更远的连接可能是以太网。

绕过 CPU：普通以太网的数据要经过 CPU（拷贝到内核的套接字缓冲区、打包成 TCP 数据包、再拷贝到网卡）。RDMA（远程直接内存访问）让一块 GPU 直接读写另一台机器上 GPU 的显存，不经过 CPU。InfiniBand 支持 RDMA；RoCE（基于融合以太网的 RDMA）让以太网也能做到，更便宜但弱一些，Meta 就用它。

新趋势：GB200 / GB300 NVL72 把 72 块 GPU（每个托盘 8 块、每个机柜 9 个托盘）放进同一个 NVLink 域，“机内高速互联”的范围从 8 块扩大到了 72 块。

2.2 TPU 的网络：环面而不是交换机

GPU 在一个 NVLink 域内是”全连接”（通过交换机，任意两块之间都能直接通信）；Google 的 TPU 则用环面网格（toroidal mesh）：每块芯片只和上下左右（3D 时还有前后）的邻居直接相连，最边上的芯片绕回到另一边。

• 网格的好处：便宜，规模可以做得很大，对结构规整的通信（比如张量并行里沿某一维的 all-reduce）很高效；
• 交换式 / 全连接的好处：对不规则的通信（比如 MoE 的 all-to-all，每个设备给每个设备发不同的数据）更友好。

据 CS336 2026 年的课件，最新一代 TPU 的网络设计也开始向交换式、树形拓扑靠拢——很可能是因为 MoE 越来越普遍。

2.3 NCCL：把”集合操作”翻译成数据包

NVIDIA 的 NCCL（NVIDIA Collective Communication Library）负责 GPU 之间的集合通信：

• 探测硬件拓扑（几个节点、几个交换机、NVLink 还是 PCIe）；
• 选择 GPU 之间的最优路径和算法（环形、树形……）；
• 启动 GPU 核函数来发送 / 接收数据。

PyTorch 的 torch.distributed 在 GPU 上调用的就是 NCCL（后端名 nccl）；在 CPU 上用 gloo。

2.4 通信版的 roofline

01_why_parallelism.py 的第三张表把各级通道的带宽放在一起，并算了一个关键的数：H100 在传输 1 字节的时间里能做多少次运算：

最后一列就是第 5 讲 roofline 里的”临界算术强度”，只是把显存带宽换成了通信带宽。读写 HBM 时，每字节配 295 次运算就能不被带宽拖慢；走 InfiniBand 时，每字节要配约 2 万次运算！

这张表决定了并行策略的基本布局：通信量大的切分方式（张量并行）只能放在机器内部（NVLink）；通信量小、或者容易和计算重叠的方式（流水线并行、数据并行）才适合跨机器。第 8 讲会反复用到这个原则。

3. 集合通信原语

3.1 为什么用”集合”通信？

最底层的通信是点对点的：设备 A 发给设备 B。但分布式训练里的通信大多有固定的”模式”：比如”每个设备都要得到所有设备的梯度之和”。集合操作（collective operations）就是把这种全局的通信模式打包成一个操作：你只说”要什么”，通信库（NCCL）负责”怎么传最快”。这比自己管理一堆点对点通信更简单，通常也更快。这些原语是上世纪 80 年代并行计算领域就有的经典概念。

术语：

• rank：设备（进程）的编号，0、1、2、……
• world size：设备总数。

3.2 八种原语

下面的例子与 02_collectives.py 的输出一致（4 个 rank）。

基础四种（结果只在一个 rank 上，或者从一个 rank 出发）：

主力三种（结果给所有人）：

最一般的一种：

块大小相同时，all-to-all 就像把一个”rank × 块”的矩阵转置；它也能处理各块大小不同的情况（MoE 里每个专家分到的 token 数不同），但块越均衡越好。

3.3 记忆口诀

CS336 课上给出的记法：

• reduce：做一个满足结合律、交换律的运算（求和、最小值、最大值）；
• scatter 是 gather 的逆；
• all：结果送到所有设备。

3.4 最重要的等式：all-reduce = reduce-scatter + all-gather

先 reduce-scatter，每个人得到”总和”的一片；再 all-gather，把各片拼成完整的总和——结果与 all-reduce 完全一样（02_collectives.py 第 8 部分验证了这一点）。而且第 4 节会看到，在带宽受限的情况下，这就是 all-reduce 最快的实现方式，两种做法的通信量相同。

为什么这个等式重要？因为把 all-reduce 拆开以后，我们可以在两步之间插入别的操作。第 8 节的 ZeRO 正是利用这一点，在”reduce-scatter 梯度”和”all-gather 参数”之间，只更新自己负责的那一片参数——从而不需要每个人都存一整份优化器状态。

4. 通信的代价：环形 all-reduce

4.1 朴素做法：所有人发给一个人

最直接的 all-reduce：所有人把向量发给 rank 0，rank 0 求和后再发回去。问题是 rank 0 要收 $(p-1)n$ 个数、发 $(p-1)n$ 个数，设备越多它越堵，而其他设备的链路大部分时间闲着。

4.2 环形算法

把 $p$ 个设备排成一个环，每个设备只和”下一个”通信。把每个设备的向量切成 $p$ 块。

• 阶段 1（reduce-scatter），共 $p-1$ 步：每一步，设备 $i$ 把一块”部分和”发给设备 $i+1$，对方把它加到自己的同一块上。$p-1$ 步之后，每个设备手里恰好有一块”完整的和”。
• 阶段 2（all-gather），共 $p-1$ 步：每一步，设备 $i$ 把一块”完整的和”发给设备 $i+1$，对方直接覆盖。$p-1$ 步之后，大家都有了全部。

03_ring_allreduce.py 用 4 个设备（设备 $r$ 的数据是 [0,1,2,3]+r）逐步打印了整个过程。截取其中几步：

初始状态：
  设备 0: [0.0] | [1.0] | [2.0] | [3.0]
  设备 1: [1.0] | [2.0] | [3.0] | [4.0]
  ...
reduce-scatter 第 3 步之后：
  设备 0: [0.0] | [10.0] | [11.0] | [9.0]     ← 第 1 块已是完整的和 10
  设备 1: [1.0] | [2.0]  | [14.0] | [13.0]    ← 第 2 块已是完整的和 14
  设备 2: [3.0] | [5.0]  | [4.0]  | [18.0]    ← 第 3 块已是完整的和 18
  设备 3: [6.0] | [9.0]  | [9.0]  | [6.0]     ← 第 0 块已是完整的和 6
all-gather 第 3 步之后：
  所有设备: [6.0] | [10.0] | [14.0] | [18.0]

阶段 1 结束时，每个设备恰好持有总和的一片——这就是一次 reduce-scatter；阶段 2 是一次 all-gather。环形 all-reduce 的实现本身就证明了 3.4 节的等式。

关键性质：每一步里，所有设备同时发一块、收一块，所有链路都在工作，没有人闲着，也没有人被堵。

4.3 通信量几乎与设备数无关

每个设备共发送 $2(p-1)$ 次，每次 $n/p$ 个数：

（向量长度 $n$ = 100 万，来自 03_ring_allreduce.py 演示 2。）

环形算法每个设备只发约 $2n$ 个数，不随设备数增长。这是后面所有通信量分析的基础：all-reduce 一个大小为 $n$ 的张量，每个设备的通信量约为 $2n$；reduce-scatter 和 all-gather 各约为 $n$。

4.4 延迟 vs 带宽：α-β 模型

发送一条消息的时间可以近似为：

$\alpha$ 是固定的延迟（启动一次通信的开销，微秒级），$\beta$ 是带宽。用它比较三种 all-reduce 算法（$\alpha$ = 5 µs，$\beta$ = 50 GB/s，03_ring_allreduce.py 演示 3）：

• 环形：$2(p-1)$ 步，每步 $n/p$ 字节：$T = 2(p-1)(\alpha + n/(p\beta))$；
• 树形（沿二叉树归约到根再广播回来）：$2\log_2 p$ 步，每步 $n$ 字节：$T = 2\log_2 p\,(\alpha + n/\beta)$；
• 朴素：$T = 2(p-1)(\alpha + n/\beta)$。

• 大消息（训练中同步梯度）：带宽项占主导，环形最快，而且几乎不随 $p$ 变慢；
• 小消息：延迟项占主导，环形要走 $2(p-1)$ 步，$p$ 大时反而最慢；树形只要 $2\log_2 p$ 步。

NCCL 会根据消息大小和拓扑自动选择算法（环形、树形等）。对我们的启示是：把很多小张量合并成大张量再通信（第 6.4 节的”分桶”），能省下大量的延迟开销。

4.5 怎么测量通信速度？

CS336 课上用 4 块 GPU 实测了 all-reduce 和 reduce-scatter。测量方法和第 6 讲的基准测试一样：先热身，每次计时前后都要同步（多卡时还要用 dist.barrier() 让所有进程对齐）。

有两种”带宽”（与 NVIDIA 的 nccl-tests 定义相同）：

• 算法带宽 algbw = 消息大小 / 耗时。直观，但会随设备数变化；
• 总线带宽 busbw = algbw × 修正系数。all-reduce 的系数是 $2(p-1)/p$，reduce-scatter 和 all-gather 是 $(p-1)/p$。修正后等于”每个设备的链路实际跑了多快”，与设备数无关，可以直接和硬件的标称带宽比较。

03_ring_allreduce.py 演示 5 在本机 CPU 上用 gloo 实测（4 个进程）：64 MB 的 all-reduce 约 24 ms，总线带宽约 4 GB/s。有一个值得注意的现象：理论上 reduce-scatter 只做 all-reduce 的”前一半”，应该快一倍；但在本机的 gloo 上，reduce-scatter 反而比 all-reduce 还慢一点（约 1.16 倍）。在 NCCL + GPU 上，二者的耗时基本符合理论。不同的通信库、不同的硬件，各个操作的实际效率可能差别很大——理论分析告诉你该期待什么，实测才能告诉你实际是什么。

5. 用 torch.distributed 编程

5.1 SPMD：同一份程序，不同的数据

几乎所有分布式训练代码都是 SPMD（Single Program, Multiple Data）风格：启动 $p$ 个进程，每个进程运行同一份代码，只是各自的 rank 不同。程序里用 rank 决定”我处理哪部分数据”、“我负责哪部分参数”。

一个最小的例子（CS336 课上的代码，稍作简化）：

import torch.distributed as dist

def main(rank, world_size):
    dist.init_process_group("gloo", rank=rank, world_size=world_size)   # 在 GPU 上用 "nccl"
    data = torch.tensor([0., 1, 2, 3]) + rank      # 每个 rank 的数据不同
    dist.all_reduce(data, op=dist.ReduceOp.SUM)     # 原地修改：所有 rank 都得到 [6, 10, 14, 18]
    dist.destroy_process_group()

5.2 本讲的代码怎么在一台电脑上模拟多卡

dist_utils.py 把样板代码封装好了：

from dist_utils import launch

def worker(rank, world_size, x):       # 每个进程都运行它
    ...
    return 结果

results = launch(worker, 4, x)         # 启动 4 个进程；results[r] 是 rank r 的返回值

launch 做的事：找一个空闲端口，用 torch.multiprocessing.spawn 启动 4 个进程，每个进程用 gloo 后端初始化进程组（大家在 127.0.0.1:端口 碰头），运行你的函数，最后把返回值通过临时文件传回主进程。

在真实集群上，通常用 torchrun --nproc_per_node=8 train.py 这样的工具启动，每台机器 8 个进程，每个进程绑定一块 GPU；多台机器时还要告诉它们主节点的地址。算法本身不需要改。

5.3 三个常见的坑

1. 集合操作不配对 → 卡死。集合操作要求所有进程都调用。如果写了 if rank == 0: dist.all_reduce(x)，其他进程永远不会调用，rank 0 就会永远等下去。这是分布式程序最常见的 bug。
2. 原地修改。dist.all_reduce(x) 直接修改 x，不返回新张量。
3. 异步操作要等待。dist.all_reduce(x, async_op=True) 立即返回一个句柄，通信在后台进行；在使用结果之前必须调用 handle.wait()。第 6.4 节会用它实现通信与计算的重叠。

6. 数据并行（DDP）

6.1 思路：每人一份模型，各算一部分数据

数据并行（Distributed Data Parallel，DDP）：

1. 每个 rank 有一份完整的模型和优化器；训练开始时从 rank 0 广播参数，保证起点相同；
2. 全局批次（global batch）切成 $p$ 份，每个 rank 只做自己那一份的前向和反向；
3. 用 all-reduce 把各 rank 的梯度求平均；
4. 每个 rank 用相同的平均梯度、相同的优化器更新，参数始终保持一致。

**为什么这是对的？**损失是对批次求平均的：

也就是说，整批的梯度 = 各 rank 梯度的平均（前提是各份大小相同）。所以数据并行在数学上完全等价于单卡用整个大批次训练。

04_data_parallel.py 演示 1 在单进程里验证了这个公式：整批梯度与 4 份梯度平均值的最大差异是 $1.9\times10^{-9}$——不是 0，只是因为浮点求和的顺序不同（第 2 讲：浮点加法不满足结合律）。

6.2 最朴素的实现

与单卡训练相比，唯一多出来的就是广播初始参数和同步梯度（CS336 课上的写法）：

for p in model.parameters():
    dist.broadcast(p.data, src=0)                   # 起点相同

for x, y in batches:
    x, y = 我的那一份(x), 我的那一份(y)
    opt.zero_grad()
    loss = F.mse_loss(model(x), y)
    loss.backward()
    for p in model.parameters():                    # ← 数据并行的全部秘密
        dist.all_reduce(p.grad, op=dist.ReduceOp.SUM)
        p.grad /= world_size
    opt.step()

04_data_parallel.py 演示 2 用 4 个进程训练 10 步，与单卡大批次训练对比：

• 各 rank 最后一步的本地损失不同（0.249、0.462、0.530、0.343），因为它们算在不同的数据上；
• 但 4 个 rank 的参数逐比特相同——梯度经过 all-reduce 后完全一样，更新也就完全一样，不需要额外同步参数；
• 与单卡训练的参数差异约 $1.6\times10^{-5}$：浮点求和顺序带来的微小梯度差异，又被 Adam 的”除以 $\sqrt{v}$“稍微放大。数学上是同一个算法。

6.3 代价分析

6.4 让通信”藏”到计算后面：重叠与分桶

朴素实现要等反向传播全部结束才开始通信，通信期间计算单元闲着。可以做得更好：

重叠。反向传播从最后一层往前算，最后一层的梯度最先算好。它一算好就可以发出去，同时继续算前面层的梯度——通信被计算”藏”起来了。PyTorch 允许给参数注册钩子：

p.register_post_accumulate_grad_hook(钩子)    # p.grad 算好时自动调用

在钩子里发起异步 all-reduce（async_op=True），把句柄存起来；optimizer.step() 之前统一 wait()。

分桶。一个模型往往有几百个参数张量，很多是很小的（LayerNorm 的权重、偏置……）。每个单独 all-reduce，就要付几百次延迟 $\alpha$（第 4.4 节）。分桶把相邻几个参数的梯度拼成一个大张量再通信：次数少了，但要等桶里所有梯度都算好才能发。桶的大小是一个权衡，PyTorch DDP 默认 25 MB。

04_data_parallel.py 的 DDPOverlap 类实现了这两种优化（这正是 CS336 作业 2 要求实现的内容）。在本机 CPU 上实测（8 层 × Linear(1024)+LayerNorm，840 万参数分布在 34 个张量里，4 个进程，每个进程每步 128 个样本）：

三种实现训练出的参数完全相同，区别只在什么时候、分几次通信。

**【边界】**这是 CPU 上的实验：gloo 的通信和矩阵乘法抢的是同一批 CPU 核，每次运行的数字也会有小的波动。在 GPU 上，计算在 GPU 上进行，通信由 NCCL 用单独的 CUDA 流和网卡完成，二者真正并行，重叠的收益通常更明显。

6.5 梯度累积：少通信几次

如果显存只够每个 rank 每次算一个小批次，可以连续做 $k$ 次前向 + 反向，把梯度累加起来，只在最后一次之后 all-reduce，再更新参数。效果等价于 $k$ 倍大的批次，而通信次数减少到 $1/k$。PyTorch DDP 用 with model.no_sync(): 跳过中间几次的同步。

6.6 数据并行的极限

批次大小的限制。每个 rank 至少要处理一个样本，所以数据并行的卡数不能超过批次大小。更现实的限制是：全局批次不能无限大。批次超过某个”临界批量大小”之后，更大的批次几乎不能再减少所需的训练步数（第 9 讲），白白浪费算力。

通信能不能被计算藏住？（05_zero_memory.py 演示 4）设参数量为 $\Psi$，每块卡每步处理 $b$ 个 token：

令计算时间 ≥ 通信时间，$\Psi$ 被约掉了：

H100、MFU 40% 时，走 InfiniBand（50 GB/s）要求每块卡每步至少约 5300 个 token，走 NVLink 约 590 个。这个结论与模型大小无关：大模型参数多、通信多，但每个 token 的计算也同比例地多。

所以：数据并行要求每块卡的批次足够大，而全局批次 = 每卡批次 × 卡数，又不能超过临界批量大小。**卡数多到一定程度后，光靠数据并行就不划算了。**这是第 8 讲模型并行的动机之一。

7. 显存：每个参数 16 字节，而且存了 N 份

数据并行的另一个大问题是显存。CS336 课件说”我们的显存状况其实糟透了”：混合精度训练时，每个参数有 5 份拷贝、共 16 字节——BF16 参数、BF16 梯度、FP32 主权重、FP32 的 $m$、FP32 的 $v$。其中后三项合称优化器状态，占了 12 字节。

而数据并行让 $N$ 块卡每块都存一整份。$N$ 块卡上存着 $N$ 份完全相同的优化器状态——这就是”冗余”。

8. ZeRO：零冗余优化器

ZeRO（Zero Redundancy Optimizer，Rajbhandari et al. 2020，微软 DeepSpeed）的核心想法：把昂贵的状态切开，每个 rank 只存 $1/N$，并利用 all-reduce = reduce-scatter + all-gather 这个等式来安排通信。

分三个阶段，逐步切得更多：

8.1 ZeRO-1：切优化器状态

做法（CS336 课件的四步）：

1. 每个 rank 在自己那份数据上算出完整的梯度；
2. reduce-scatter 梯度：每个 rank 只得到自己负责的那一片参数的梯度之和（通信约 $\Psi$）；
3. 每个 rank 用这片梯度和自己保存的那片优化器状态，只更新自己那片参数；
4. all-gather 参数：把各片更新后的参数拼回完整模型（通信约 $\Psi$）。

**为什么说它是”免费的”？**DDP 本来就要做一次 all-reduce（$2\Psi$）；ZeRO-1 只是把它拆成 reduce-scatter（$\Psi$）+ all-gather（$\Psi$），在中间插入”只更新自己那一片”。**通信量一点没变，优化器状态却只剩 $1/N$。**按 CS336 课件的说法：既然是免费的，不如一直开着。

06_zero1_sharded_optimizer.py 给出两种实现：

• ShardedOptimizer（作业 2 的接口）：按”整个张量”把参数分给各 rank，每个 rank 的内部优化器只管自己的参数；梯度照常 all-reduce，更新后每个参数由它的”主人”广播给大家。结果与 DDP 逐比特相同，每个 rank 的 AdamW 状态约为不分片时的 1/4（33 KB vs 134 KB）。它简单、容易接入，但通信量是 all-reduce（$2\Psi$）+ 广播（$\Psi$）= $3\Psi$。
• FlatZeRO1（ZeRO 论文的做法）：把所有参数拼成一个长向量，均匀切成 $N$ 段；梯度 reduce-scatter，更新自己那段，参数 all-gather。通信量 $2\Psi$，切得也完全均匀。与 DDP 的差异约 $10^{-6}$（reduce-scatter 和 all-reduce 内部的求和顺序不同）。

8.2 ZeRO-2：再切梯度

既然每个 rank 只需要自己那一片梯度，能不能连完整的梯度都不存？

难点：数据并行中，每个 rank 必须在自己的数据上算出完整的梯度——反向传播就是这样算的。ZeRO-2 的办法是边算边送、送完就扔：

1. 反向传播逐层进行；某一层的梯度一算好，就立即 reduce（发给负责它的 rank）；
2. 这一层的梯度一旦不再被反向传播需要，就立即释放；
3. 每个 rank 用自己那片梯度更新自己那片参数；
4. all-gather 参数。

通信量仍是约 $2\Psi$（几乎免费，只是有一些额外开销）。代价是实现更复杂：它必须和反向传播紧密配合，类似第 6.4 节的钩子。

8.3 ZeRO-3 / FSDP：连参数也切

最后一步：参数也只存 $1/N$。可前向计算需要完整的参数啊？答案是用的时候再拼，用完马上扔：

前向，第 ℓ 层：all-gather 这层的参数 → 计算 → 释放完整参数（只留自己那片）
反向，第 ℓ 层：再 all-gather 一次这层的参数 → 算输入梯度和权重梯度 → 释放完整参数
               → reduce-scatter 权重梯度（每人只拿自己那片的梯度之和）
更新：每个 rank 只更新自己那片（和 ZeRO-1 一样）

PyTorch 的 FSDP（Fully Sharded Data Parallel）就是 ZeRO-3 的实现。

通信量：前向 all-gather $\Psi$ + 反向 all-gather $\Psi$ + reduce-scatter $\Psi$ = $3\Psi$，比 DDP 多 50%。CS336 课件的评价：“不算太糟”——换来的是模型状态完全均摊。

能不能不被通信拖慢？关键是重叠：计算第 $\ell$ 层的同时，就提前 all-gather 第 $\ell+1$ 层的参数（预取）。只要每层的计算时间比 all-gather 一层参数的时间长，通信就被完全藏住了。代价是同时存在两层的完整参数。

07_fsdp.py 用 torch.autograd.Function 实现了一个”分片线性层”：

class ShardedLinearFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, shard, ...):
        w = all_gather(shard)                  # 拼出完整权重
        ctx.save_for_backward(x, shard)        # 注意：只保存分片，不保存完整权重
        return x @ w.T                         # 返回后 w 不再被引用 → 立即释放

    @staticmethod
    def backward(ctx, dy):
        x, shard = ctx.saved_tensors
        w = all_gather(shard)                  # 反向时再拼一次
        dx = dy @ w
        dw = dy.T @ x
        grad_shard = reduce_scatter(dw) / N    # 每人只拿自己那片梯度
        return dx, grad_shard, ...

（torch.autograd.Function 让你自己定义前向和反向怎么算：forward 里的运算不会被自动求导记录，backward 返回每个输入的梯度。）

实测结果（4 个进程，5 层共 16448 个参数）：

• 每个 rank 常驻的参数只有 4112 个（$\approx\Psi/4$），优化器也只看得到这一片；
• 训练结果与单卡大批次训练的差异 $3.9\times10^{-6}$；
• 每步 all-gather $2\Psi$、reduce-scatter $\Psi$，合计 $3\Psi$，与公式一致；
• 同一时刻最多只存在 1 个完整的权重矩阵（代码用 weakref.finalize 追踪完整参数何时被释放）。

8.4 显存对比

ZeRO 论文图 1 的例子（05_zero_memory.py 演示 1）：75 亿参数，64 块卡，混合精度 Adam（$K$ = 12）：

CS336 课件的例子（演示 2）：8 块 80 GB 的 A100，除主权重外全用 BF16（12 字节/参数）：

通信量（每步、每块卡，05_zero_memory.py 演示 3）：DDP、ZeRO-1、ZeRO-2 都是约 $2\Psi$，ZeRO-3 约 $3\Psi$。

8.5 ZeRO 解决不了的问题

1. 激活值。ZeRO 切的是模型状态，不切激活值：每个 rank 依然要存自己那份数据的全部激活值。大模型、长序列时，激活值可能比模型状态还大。
2. 批次大小。ZeRO 仍然是数据并行，第 6.6 节的限制依然存在：卡数不能超过批次大小，全局批次也不能超过临界批量大小。
3. ZeRO-3 的通信是”参数大小”的。每一步都要把整个模型的参数传两遍。模型非常大、网络又不够快时，这很贵。

这些问题需要换一种切法：切模型本身，传激活值而不是传参数——这就是下一讲的模型并行。

8.6 实践中怎么用

CS336 课件整理了一些模型的做法：Dolma / OLMo 的 7B 模型用 FSDP；DeepSeek 系列用 ZeRO-1 配合其他并行方式（DeepSeek-V3 用 ZeRO-1 + 16 路流水线并行 + 64 路专家并行）；Gemma 2 用 ZeRO-3 + 模型并行 + 数据并行。经验上：ZeRO-1 几乎总是值得开；模型不太大、单机或少数几台机器能放下时，FSDP 是最省事的选择（只要包一层，不用改模型代码）。

9. 本讲小结

1. 为什么多卡：显存放不下（每参数 16 字节），算力不够（6ND）。目标是显存和算力都随卡数线性扩展。
2. 层次结构：HBM → NVLink → InfiniBand → 以太网，每往外一级带宽低一个数量级。走 InfiniBand 时，每传 1 字节要配约 2 万次运算才不会被拖慢。
3. 集合通信：broadcast、scatter、gather、reduce、all-gather、reduce-scatter、all-reduce、all-to-all。all-reduce = reduce-scatter + all-gather。
4. 环形 all-reduce：每个设备发送 $2(p-1)/p \cdot n \approx 2n$，与设备数几乎无关。大消息看带宽（环形好），小消息看延迟（树形好）。
5. DDP：广播初始参数 + all-reduce 平均梯度，数学上等价于大批次训练。通信 $2\Psi$，显存不省。重叠与分桶能把通信藏到反向传播后面。
6. ZeRO：ZeRO-1 切优化器状态（通信不变，“免费”）；ZeRO-2 再切梯度（几乎免费）；ZeRO-3 / FSDP 再切参数（通信 $3\Psi$，用时拼、用完扔、预取重叠）。
7. 数据并行的极限：不切激活值；受批次大小限制；每块卡每步需要几千个 token 才能藏住跨机通信。

10. 与作业 2 的关系

作业 2（系统）的后半部分是分布式训练，与本讲直接对应：

• 通信基准测试：在不同后端（gloo / nccl）、不同进程数、不同数据大小下测量 all-reduce 的耗时（第 4.5 节、03_ring_allreduce.py）；
• 朴素 DDP：广播参数 + 反向后逐个 all-reduce 梯度，并测量通信占总时间的比例（第 6.2 节）；
• 重叠通信的 DDP：用梯度钩子和异步 all-reduce，先按单个参数、再按桶（第 6.4 节的 DDPOverlap）；
• 优化器状态分片：实现 ShardedOptimizer（第 8.1 节），比较显存和速度；
• 可能还有关于 FSDP、多种并行组合的资源核算问题（第 8 讲会讲张量并行、流水线并行）。

建议：先在 CPU 上用 gloo 把正确性调通（与单卡训练对比参数），再到 GPU 上用 nccl 测速度。

11. 练习

练习 1（集合操作）：4 个 rank 的输入分别是 [1,2,3,4]、[5,6,7,8]、[9,10,11,12]、[13,14,15,16]。写出下列操作之后每个 rank 的结果：(a) all-reduce（求和）；(b) reduce-scatter（求和）；(c) 对 (b) 的结果做 all-gather；(d) all-to-all。(e) 如果每个 rank 的输入只有 2 个元素，上面哪些操作还能直接做？

(a) 所有人 [28, 32, 36, 40]。(b) rank 0 [28]，rank 1 [32]，rank 2 [36]，rank 3 [40]。(c) 所有人 [28, 32, 36, 40]，与 (a) 相同。(d) rank $j$ 收到每个人的第 $j$ 个元素：rank 0 [1,5,9,13]，rank 1 [2,6,10,14]，rank 2 [3,7,11,15]，rank 3 [4,8,12,16]。(e) all-reduce 可以；reduce-scatter 和（均分的）all-to-all 要求长度是 rank 数的整数倍，长度 2 不能均分给 4 个 rank。实际代码中常常要补零到能整除（FlatZeRO1 和 07_fsdp.py 都这样做了）。

练习 2（通信量）：8 块卡做数据并行，模型 70 亿参数，梯度用 BF16。(a) 用环形 all-reduce，每块卡每步发送多少字节？(b) 如果链路带宽是 50 GB/s，通信至少要多久？(c) 如果改用 ZeRO-3 呢？

(a) $2 \times 7/8 \times 7\times10^9 \times 2$ 字节 ≈ 24.5 GB。(b) 24.5 GB / 50 GB/s ≈ 0.49 秒。(c) $3 \times 7/8 \times 14$ GB ≈ 36.8 GB，约 0.74 秒。

练习 3（α-β 模型）：对 $p$ 个设备，环形和树形 all-reduce 的耗时相等时，消息大小 $n$ 是多少？用 $\alpha$ = 5 µs、$\beta$ = 50 GB/s、$p$ = 64 计算。

令 $2(p-1)(\alpha + n/(p\beta)) = 2\log_2 p\,(\alpha + n/\beta)$。代入 $p$ = 64：$63\alpha + 63n/(64\beta) = 6\alpha + 6n/\beta$，即 $57\alpha = n(6 - 63/64)/\beta ≈ 5.016\,n/\beta$，$n ≈ 57 \times 5\times10^{-6} \times 50\times10^9 / 5.016 ≈ 2.8$ MB。小于约 3 MB 时树形更快，大于时环形更快。

练习 4（等价性）：数据并行中，如果各 rank 的局部批次大小不同（比如 rank 0 有 20 个样本，其余各 12 个），直接对各 rank 的梯度求平均还等价于大批次训练吗？怎么修正？

不等价。rank 0 的每个样本的权重变小了。修正：每个 rank 先把梯度乘以自己的样本数 $b_r$（即变成”梯度之和”），all-reduce 求和后再除以总样本数 $\sum_r b_r$。或者直接让每个 rank 计算损失之和（而不是平均），最后除以全局批次大小。

练习 5（为什么参数不需要同步）：朴素 DDP 只同步梯度，从不同步参数。为什么各 rank 的参数能一直保持相同？什么情况下会”跑偏”？

初始参数相同（广播），每步的平均梯度相同（all-reduce），优化器是确定性的，所以更新后的参数相同，归纳下去永远相同。会跑偏的情况：初始化时忘了广播；优化器里有随机性（比如某些 dropout 或随机舍入的实现依赖各 rank 不同的随机种子）；各 rank 的数值计算不确定（例如 GPU 上某些非确定性的核函数，使梯度在最后几位上不同——all-reduce 后仍然相同，但如果某个 rank 用了本地梯度做别的事就会不同）；以及硬件故障。

练习 6（ZeRO 显存）：用 64 块 80 GB 的卡训练一个 1000 亿参数的模型，混合精度 Adam（16 字节/参数）。DDP、ZeRO-1、ZeRO-2、ZeRO-3 各需要每块卡多少显存存放模型状态？哪些能放下？

DDP：1.6 TB；ZeRO-1：$4\times100 + 12\times100/64$ = 418.75 GB；ZeRO-2：$2\times100 + 14\times100/64$ ≈ 221.9 GB；ZeRO-3：$1600/64$ = 25 GB。只有 ZeRO-3 能放下（还要留出激活值的空间）。

练习 7（为什么 ZeRO-3 要 all-gather 两次）：ZeRO-3 在前向时已经 all-gather 过一次参数了，反向时为什么还要再来一次？有没有办法只做一次？代价是什么？

因为前向结束后完整参数被释放了（这正是 ZeRO-3 省显存的方式），反向算 $dx = dy\,W$ 时又需要完整的 $W$。只做一次的办法是前向后不释放，一直保留到反向——但这样所有层的完整参数会同时存在，显存和 DDP 一样，ZeRO-3 就白做了。折中：只对最后几层保留（它们马上就要做反向），PyTorch FSDP 的 reshard_after_forward 选项就是控制这个的。这又是”存下来 vs 再通信一次”的权衡。

练习 8（通信能否被藏住）★：在 6.6 节的推导中，如果用 ZeRO-3（通信 $3\Psi$），并把梯度通信从 BF16 改成 FP32（4 字节），每块卡每步至少需要多少 token 才能藏住跨机通信？这对全局批次有什么影响（假设 1024 块卡）？

ZeRO-3 的通信是 all-gather 参数 $2\Psi$（BF16，2 字节）+ reduce-scatter 梯度 $\Psi$（FP32，4 字节）= $8\Psi$ 字节。$b \ge 8 \times 3.96\times10^{14} / (6 \times 5\times10^{10}) ≈ 10560$ 个 token。1024 块卡的全局批次至少约 1080 万 token。Llama 3 405B 的全局批次是 1600 万 token 左右（训练后期），同一数量级——这也是为什么大规模训练的批次都这么大。

练习 9（动手）： (a) 在 04_data_parallel.py 的 naive_ddp_worker 里加入梯度累积（2 个微批次、只同步一次），验证结果不变。 (b) 修改 DDPOverlap 的桶大小（1 MB、64 MB），记录每步耗时和通信次数，画出”桶大小 → 耗时”的曲线。 (c) 在 07_fsdp.py 里实现预取：第 $\ell$ 层前向时用 async_op=True 发起第 $\ell+1$ 层的 all-gather。peak_alive 会变成几？

12. 延伸阅读

• CS336 Spring 2026 Lecture 7（GitHub stanford-cs336/lectures 中的 lecture_07.py，可执行讲义，包含本讲的集合操作、DDP 以及张量 / 流水线并行的最简实现）和 Lecture 8（lecture_08.pdf）。
• Rajbhandari et al.，ZeRO: Memory Optimizations Toward Training Trillion Parameter Models（2020）：ZeRO-1/2/3 的原始论文。
• Zhao et al.，PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel（2023）：FSDP 的设计与工程细节（预取、按块分片、混合精度）。
• Li et al.，PyTorch Distributed: Experiences on Accelerating Data Parallel Training（2020）：PyTorch DDP 的梯度分桶与重叠。
• NVIDIA nccl-tests 的 PERFORMANCE.md：算法带宽与总线带宽的定义。
• Stas Bekman，Machine Learning Engineering（GitHub stas00/ml-engineering）：网络、通信基准测试、实际训练中的各种坑。
• Google，How to Scale Your Model（“Scaling Book”，jax-ml.github.io/scaling-book）：从 TPU 的角度系统讲解各种并行方式的通信代价。

13. 下一讲预告

数据并行（加上 ZeRO）解决了”模型状态放不下”的问题，但留下了两个难题：激活值没有被切分；批次大小限制了能用多少块卡。下一讲（并行 II）换一种思路——切模型本身：

• 流水线并行：按层切，每块卡负责几层，激活值在卡之间传递。怎么减少”流水线气泡”？
• 张量并行：把每一层的矩阵切开（Megatron 式按列 / 按行切分），为什么只能在机器内部用？
• 序列并行 / 上下文并行：把激活值沿序列长度切开，环形注意力怎么用上第 5 讲的在线 softmax？
• 专家并行：MoE 的专家放在不同的卡上，all-to-all 派发 token；
• 怎么把它们组合成 3D / 4D 并行，以及 Llama 3、DeepSeek-V3 等模型的实际配置。

下一讲之前，建议你：跑通 code/lecture07/ 下的 7 个脚本，读懂 dist_utils.py、DDPOverlap 和 ShardedLinearFn，完成练习 2、4、6、7。