ZeRO-3：省下副本之后，怎样算完一层？

先锁定一个不会变的条件：四张卡各有自己的数据，每一份数据仍要经过模型的全部层。我们现在跟踪其中一层，并不把第一层永久交给第一张卡、第二层永久交给第二张卡。图上的四种颜色表示这个层的参数分片归谁保存。

轮到这一层时，四张卡交换它们负责的分片，于是每张卡都临时拥有这一层的完整参数。这叫 all-gather，中文可以直接理解为收集到完整视图。箭头送来的是本层需要的参数，不要求一次把全模型所有层都聚齐。

现在材料齐了，各卡就可以用完整本层参数处理自己的输入，产生交给下一层的中间结果。我们把这些中间结果叫激活。然后按同样办法继续后面的层，最终形成各卡数据对应的损失。参数相同不意味着激活相同，因为各卡读的数据不同。

本例采用最容易看清节省来源的策略：本层前向算完，就释放刚才临时收齐的完整参数，只保留自己负责的分片。因此全模型没有在每卡长期变回完整副本。代价也已经埋下：一会儿反向如果还需要这些参数，我们要再收集一次。

前向走完后，从损失沿层往回传播。轮到这一层时，再次收集本层参数，用它和反向传来的信息计算梯度。这里画了第二次收集，是为了让通信代价真实可见。工程上可以预取或者延迟释放，但那又会改变峰值内存；没有一种图能同时免费省掉所有存储与通信。

四张卡各算出了本地数据的贡献。接下来先把对应梯度汇总，再把合并后的各段发到负责那一段参数的卡。这叫 reduce-scatter：归约后分片。它不是随便分走四份未经合并的梯度；每个最终分片都包含四份数据的贡献，平均方式仍遵守训练的约定。

最后各卡只更新自己负责的主权重和动量，并得到新的工作参数分片。下一轮收集时，拿到的就是更新后的参数。这一圈走完，节省从哪里来、通信为什么出现，就连在一起了。DeepSpeed ZeRO、PyTorch 的全分片数据并行 FSDP，是可以使用的工程路线；具体预取和驻留不同，不能当成逐细节相同的软件。接下来我们换一种切法：直接把计算分给不同设备。