起因#
之前几篇文章从模型显存组成一路算到 7B 模型的训练账本，结论很直接：BF16 mixed precision + AdamW 的 Model States 就有 112 GB，一张 80 GB A100 连放都放不下。Checkpointing + Offloading 两个一起用可以勉强塞进去，但训练速度要被 PCIe 传输和额外重计算拖累。
那篇结尾我写了一句"下一步自然是考虑让多张 GPU 分担训练状态"。后来开始看多卡训练的资料，发现第一个问题不是"怎么拆 model states"，而是更前面的事情：多张 GPU 各自算出来的梯度数值不一样，想保持模型参数一致，数据应该怎么在 GPU 之间交换？
我最开始的直觉是：全部发给一个中心节点求平均不就行了？后来发现，实际用的通信方式比"收集到一起再广播回去"要精巧得多——不同的通信原语决定了每张 GPU 最终持有什么数据，也直接影响了 DDP 和 ZeRO 在显存上的差异。
这篇先把三个最核心的通信原语搞清楚，再看 Ring 算法是怎么用邻居通信替代全局交换的。
两张 GPU 各算各的，梯度就是不一样#
用最简单的例子。两张 GPU，模型参数完全相同，各拿到一个不同的 mini-batch 做 forward 和 backward。
Backward 结束后，每张 GPU 都算出了一份完整的梯度——覆盖模型的所有可训练参数，尺寸完全一样。但因为训练数据不同，这两份梯度的数值通常不同。
这两件事不矛盾，但我一开始容易混着理解：
- “完整"说的是覆盖范围：每张 GPU 对所有参数都有梯度，不是只算了一部分参数的梯度。
- “不同"说的是数值：同一个参数在两张 GPU 上的梯度值大概率不一样，因为输入数据不同。
如果两张 GPU 拿着不同的梯度各自更新参数，一步之后参数就不一致了。下一步的 forward 结果也会不同，模型从此分叉。
所以在参数更新之前，需要把梯度同步——让所有 GPU 对每个参数达成一致的梯度值，最常见的做法是求平均。
问题来了：多个 GPU 之间的数据交换，有哪些不同的方式？
三种交换方式，结果很不一样#
先用具体数字看清楚区别。假设有 2 张 GPU，各自有一段数据：
- GPU 0：
[1, 2, 3, 4] - GPU 1：
[5, 6, 7, 8]
AllGather：收集，不计算#
每张 GPU 把自己的数据发给所有人。操作结束后，每张 GPU 都拿到完整的数据集合：
- GPU 0：
[1, 2, 3, 4, 5, 6, 7, 8] - GPU 1：
[1, 2, 3, 4, 5, 6, 7, 8]
没有做任何数学运算。只是把分散在不同 GPU 上的数据拼到一起，每张 GPU 各一份完整副本。
AllReduce：归约，每卡得到相同的完整结果#
对各 GPU 对应位置的数据做归约（这里用求和），所有 GPU 最终拿到相同的完整归约结果：
- GPU 0：
[1+5, 2+6, 3+7, 4+8]=[6, 8, 10, 12] - GPU 1：
[1+5, 2+6, 3+7, 4+8]=[6, 8, 10, 12]
每个位置都做了加法，两张 GPU 结果完全一致。
这正是数据并行训练同步梯度最直接的需求：每张 GPU 都拿到所有 GPU 的梯度之和（除以 GPU 数量就是平均值），然后各自用相同的梯度更新参数。
ReduceScatter：归约，但每卡只拿一部分#
先归约，再分片。每张 GPU 最终只拿到归约结果的一个分片：
- 归约结果（求和）：
[6, 8, 10, 12] - GPU 0 拿到前半：
[6, 8] - GPU 1 拿到后半：
[10, 12]
和 AllReduce 相比：归约是一样的，但 ReduceScatter 不要求每张 GPU 都持有完整结果。每张 GPU 只负责保存归约结果的一个分片。
第一次看到 ReduceScatter 的时候，我不太理解为什么需要这种"归约了但只拿一部分"的操作——既然费力气做了归约，为什么不让每卡都拿到完整结果？后来看 ZeRO 的设计才明白：如果每卡只负责更新一部分参数，那它只需要那部分参数对应的归约梯度就够了，拿完整结果反而浪费显存。
对比#
| 操作 | 是否归约 | 每卡最终持有 |
|---|---|---|
| AllGather | 否 | 完整数据（所有 GPU 的拼接） |
| AllReduce | 是 | 完整归约结果（所有 GPU 相同） |
| ReduceScatter | 是 | 归约结果的一个分片 |
三种操作的输入输出对比：
到这里，三种操作只定义了"输入什么、输出什么”，没有说数据具体怎么在 GPU 之间传输。最朴素的做法是每对 GPU 直接互发——但 GPU 数量一多，这种全对全通信的开销会很大。
四张 GPU 排成一个环#
Ring 算法的想法是：不需要每张 GPU 和所有 GPU 直接通信。把 GPU 排成一个逻辑环，每张 GPU 只和左右邻居交换数据，通过多轮接力传递，实现全局的数据收集或归约。
先从 Ring AllGather 看起，因为它不涉及归约，只是收集数据，比较直观。
Ring AllGather：接力传递#
4 张 GPU 各有一块数据：
- GPU 0：A
- GPU 1：B
- GPU 2：C
- GPU 3：D
组成逻辑环：GPU 0 → GPU 1 → GPU 2 → GPU 3 → GPU 0。每张 GPU 只能向右边的邻居发送数据。
第一轮：每张 GPU 发送自己的数据给右邻居。
| GPU | 发送 | 接收 | 轮后持有 |
|---|---|---|---|
| GPU 0 | A → GPU 1 | D ← GPU 3 | A, D |
| GPU 1 | B → GPU 2 | A ← GPU 0 | B, A |
| GPU 2 | C → GPU 3 | B ← GPU 1 | C, B |
| GPU 3 | D → GPU 0 | C ← GPU 2 | D, C |
第二轮：每张 GPU 转发上一轮收到的数据给右邻居。
| GPU | 发送 | 接收 | 轮后持有 |
|---|---|---|---|
| GPU 0 | D → GPU 1 | C ← GPU 3 | A, D, C |
| GPU 1 | A → GPU 2 | D ← GPU 0 | B, A, D |
| GPU 2 | B → GPU 3 | A ← GPU 1 | C, B, A |
| GPU 3 | C → GPU 0 | B ← GPU 2 | D, C, B |
第三轮：继续转发上一轮收到的。
| GPU | 发送 | 接收 | 轮后持有 |
|---|---|---|---|
| GPU 0 | C → GPU 1 | B ← GPU 3 | A, B, C, D |
| GPU 1 | D → GPU 2 | C ← GPU 0 | A, B, C, D |
| GPU 2 | A → GPU 3 | D ← GPU 1 | A, B, C, D |
| GPU 3 | B → GPU 0 | A ← GPU 2 | A, B, C, D |
三轮结束，4 张 GPU 都拿到了 A、B、C、D。AllGather 完成。
这里有一个让我印象比较深的地方：每张 GPU 不仅是自己数据的发送者，还充当了其他 GPU 数据的中转站。 比如 GPU 2 在第三轮转发了 A——A 本来是 GPU 0 的数据，GPU 0 和 GPU 2 之间没有直接通信，A 是经过 GPU 0 → GPU 1 → GPU 2 中转过来的。
4 张 GPU 用了 3 轮，也就是 \(N - 1\) 轮。每张 GPU 每轮只发送一块数据，整个过程中一共发送 \(N - 1\) 次。
Ring ReduceScatter：传递的同时做加法#
Ring AllGather 只是传递和收集。Ring ReduceScatter 不一样——数据在沿环传递的过程中被累加。
还是 4 张 GPU。这次每张 GPU 持有完整的梯度向量，把它切成 4 个 chunk：
- GPU 0：\(a_0, a_1, a_2, a_3\)
- GPU 1：\(b_0, b_1, b_2, b_3\)
- GPU 2：\(c_0, c_1, c_2, c_3\)
- GPU 3：\(d_0, d_1, d_2, d_3\)
目标是每张 GPU 最终拿到一个 chunk 的完整归约值（4 张 GPU 对应位置之和）。
用一个 chunk 的视角追踪它怎么沿环走。以 chunk 0 为例：
第一轮：GPU 0 发送 \(a_0\) 给 GPU 1。GPU 1 收到后累加自己的 chunk 0：\(a_0 + b_0\)。
第二轮：GPU 1 把累加结果 \(a_0 + b_0\) 发送给 GPU 2。GPU 2 累加：\(a_0 + b_0 + c_0\)。
第三轮：GPU 2 把 \(a_0 + b_0 + c_0\) 发送给 GPU 3。GPU 3 累加：\(a_0 + b_0 + c_0 + d_0\)。
三轮之后，GPU 3 持有 chunk 0 的完整归约值。
同时，chunk 1 从 GPU 1 出发走同样的路线，chunk 2 从 GPU 2 出发，chunk 3 从 GPU 3 出发。四个 chunk 在同一时间并行推进，每轮每张 GPU 发送一个 chunk、接收一个 chunk，通信带宽没有浪费。
三轮结束后：
| GPU | 持有的完整归约 chunk |
|---|---|
| GPU 0 | \(a_1 + b_1 + c_1 + d_1\)（chunk 1） |
| GPU 1 | \(a_2 + b_2 + c_2 + d_2\)（chunk 2） |
| GPU 2 | \(a_3 + b_3 + c_3 + d_3\)（chunk 3） |
| GPU 3 | \(a_0 + b_0 + c_0 + d_0\)（chunk 0） |
每张 GPU 只持有归约结果的 \(\frac{1}{4}\)，但那一份是完全归约好的——所有 GPU 对应位置的梯度之和。
回过头对比两种 Ring 操作：
- Ring AllGather：数据沿环传递，接收方存储，不做运算。最终每卡持有全部数据。
- Ring ReduceScatter：数据沿环传递，接收方累加到自己的对应 chunk。最终每卡只持有一个完整归约的分片。
两者的区别就在于接收方有没有做加法。 逻辑环的结构、传递方向、轮次数量完全一样。
Ring AllReduce 是两个阶段#
到这里就有一个很自然的想法：ReduceScatter 让每张 GPU 拿到一个归约好的分片，AllGather 让每张 GPU 收集所有分片——把两步连起来，每张 GPU 就拿到了完整的归约结果。
$$ \boxed{\text{Ring AllReduce} = \text{ReduceScatter} + \text{AllGather}} $$第一阶段 ReduceScatter：通过 \(N - 1\) 轮 Ring 传递与累加，每张 GPU 得到完整归约结果的一个分片。
第二阶段 AllGather：再通过 \(N - 1\) 轮 Ring 传递与收集，每张 GPU 拿到所有分片，拼成完整的归约结果。
通信量#
用简单的数字估一下每张 GPU 的发送量。4 张 GPU，每张 GPU 的梯度大小 4 GB。梯度被切成 4 个 chunk，每个 1 GB。
ReduceScatter 阶段：每张 GPU 每轮发送一个 1 GB 的 chunk，共 \(N - 1 = 3\) 轮。单 GPU 累计发送 3 GB。
AllGather 阶段：同样每轮发一个 chunk，3 轮。单 GPU 累计发送 3 GB。
Ring AllReduce 总计：每张 GPU 累计发送 \(3 + 3 = 6\) GB。
这里统计的是单张 GPU 的发送量，接收量另算（也是 6 GB）。推广到 \(N\) 张 GPU、每张梯度 \(M\) bytes：
$$ \text{单 GPU 发送量} = 2 \times \frac{N-1}{N} \times M $$需要注意：这是 Ring 这一种特定实现下的通信量。 AllReduce 不是必须用 Ring 实现——还有 Tree AllReduce、Recursive Halving-Doubling 等其他算法，通信量和延迟特性各有不同。但这些不在这篇的范围里。
DDP 和 ZeRO Stage 2 各自选了什么#
前面一直在讲通信原语的定义和 Ring 的实现方式，最后落回到工程上。
DDP：AllReduce 同步梯度#
数据并行（DDP）的做法比较直接：每张 GPU 拥有完整的模型副本，各自用不同的 mini-batch 做 forward 和 backward。Backward 结束后通过 AllReduce 同步梯度，所有 GPU 拿到相同的梯度平均值，然后各自执行相同的 optimizer step，参数保持一致。
从通信原语看，DDP 需要的是 AllReduce——每张 GPU 既参与归约，又拿到完整的归约结果。
ZeRO Stage 2：ReduceScatter + AllGather#
ZeRO Stage 2 做了一个调整：不需要每张 GPU 都持有完整的归约梯度。
既然 optimizer step 可以按参数分组，让每张 GPU 只负责更新一部分参数就行了。那么每张 GPU 只需要拿到自己负责的那部分参数对应的归约梯度——ReduceScatter 正好做这件事。
参数更新完成后，每张 GPU 持有的模型参数不完整（各自只更新了自己负责的部分）。这时候再用 AllGather 把更新后的参数同步到所有 GPU，使每卡恢复完整的模型副本。
用通信原语描述：
- Backward 后：梯度 ReduceScatter（每卡拿到自己负责的梯度分片）
- Optimizer step：每卡用自己的梯度分片更新对应的参数
- Step 后：参数 AllGather（所有卡拿到完整的更新后参数）
通信量比 DDP 少了吗？#
这个问题我想了一会儿。乍一看，ReduceScatter 比 AllReduce"少”——毕竟每卡只拿到一个分片。但参数更新后又做了一次 AllGather。
在标准 Ring 实现下，如果梯度和参数的字节数相同：
- DDP 的 AllReduce = ReduceScatter + AllGather（Ring AllReduce 本身就拆成这两步）
- ZeRO Stage 2 = 梯度 ReduceScatter + 参数 AllGather
两者的总通信字节数可以非常接近。
ZeRO Stage 2 的核心收益不是减少通信量，而是减少每张 GPU 上梯度和 optimizer states 的冗余存储。 DDP 下每张 GPU 都存一份完整梯度和完整 optimizer states。ZeRO Stage 2 下，梯度在 ReduceScatter 后只保留自己负责的分片，对应的 optimizer states 也只维护这个分片——其他分片不归自己管，不用存。这就是之前讨论的 model states 开销能被削减的原因。
还有一个容易被忽略的细节：ZeRO Stage 2 的 backward 仍然在计算完整模型的梯度——它不是只算"自己负责的参数"的梯度。链式法则要求 backward 从 loss 一路传播到所有参数，这个计算过程无法按参数分片跳过。但梯度可以随 bucket 逐步归约和释放：一个 bucket 的梯度 ReduceScatter 完成后，不属于自己的分片可以立即丢掉，不需要整份本地梯度同时常驻显存。
写在最后#
走完这一圈，有三件事让我觉得比较重要。
通信原语不只是"数据怎么交换"，更决定了最终每张 GPU 需要持有什么数据。 AllReduce 意味着每卡都存一份完整归约结果；ReduceScatter 意味着每卡只存一个分片。选择哪种原语，直接影响显存布局。
Ring 的关键不是"环形拓扑"本身，而是"分块 + 邻居通信 + 逐轮转发"这套机制。 每张 GPU 不需要和所有 GPU 直接交换数据，通过邻居中转就能在 \(N - 1\) 轮内完成全局的收集或归约。AllReduce 可以拆成 ReduceScatter 和 AllGather 两个阶段，每个阶段都用同样的 Ring 机制实现。
DDP 和 ZeRO Stage 2 使用不同的通信原语组合，根本原因是参数更新的责任分配不同。 DDP 让每卡都更新全部参数，需要 AllReduce 拿到完整梯度。ZeRO Stage 2 让每卡只更新一部分参数，只需要 ReduceScatter 拿到对应的梯度分片——省的是存储，不必然省通信。
通信原语搞清楚以后，下一步可以更具体地看 ZeRO 的显存分片策略——Stage 2 和 Stage 3 分别在 Model States 的哪些部分做分片，分片的粒度是什么，参数什么时候需要 AllGather 回来。
参考资料#
- 7B 模型只有 14 GB 权重，一张 80 GB GPU 为什么仍可能 OOM？
- GPU 显存不够怎么办？Checkpointing 与 Offloading 的工程权衡
- 从参数量到 Peak Memory：模型显存到底应该怎么算？
- 从 Forward 到 Optimizer：再看一遍 Training Loop
- Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020.
- Patarasuk & Yuan. Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations. JPDC 2009.
- MIT 6.S191: Introduction to Deep Learning — Lecture 9: Secrets to Massively Parallel Training. https://introtodeeplearning.com/