正文
图2 混合精度训练示意图
现在，可以来计算模型在训练时需要的存储大小了，假设模型的参数大小是 Φ ，以byte为单位，存储如下（这里暂不将activation纳入统计范围内）：
在“必要存储”这一类别中，包含了模型的参数、优化过程中的动量以及方差，这些数据均采用32位浮点数（即fp32格式）进行存储，每项数据占用的空间为4Φ，整体累计占用空间为12Φ。“临时存储”类别则涵盖了以16位浮点数（即fp16格式）存储的模型参数和计算过程中产生的梯度信息，每项数据占用的空间为2Φ，累计占用空间为4Φ。综合所有类别，总的存储空间需求为16Φ。需要指出的是，因为我们一直使用Adam优化方法，因此涉及到了动量和方差的存储。如果选用其他的优化算法，这些数据可能就不需要存储。为了表述的普适性，必须将模型存储的数据量定义为KΦ，这样，总的内存消耗可以表示为2Φ+2Φ+KΦ。
在了解了哪些元素占用存储空间以及它们各自占用的存储大小之后，我们便可以探讨如何优化存储使用。可以观察到，在训练过程中，并非所有数据都需要持续保留。例如：
基于这些观察，ZeRO优化技术采取了一种直接的方法：如果某些数据在计算完成后就不再需要，那么在需要时再从其他来源重新获取这些数据，从而节省存储空间。下面我们就可以讨论3个阶段是如何通过减少不必要的数据保留，优化了存储资源的使用。
Stage 1 优化器参数划分，每个节点仅更新自己分片的参数。
如图3所示，整体流程如下：
图3 优化器优化示意图
(1) 每个GPU保存完整的模型参数W。将一个batch的数据分成n份，分给各个GPU。做完一轮的前向和反向计算后，每个GPU会得到一份梯度。
(2) 对梯度做一次环状全归约，得到完整的梯度，对All-Reduce不熟悉的朋友可以去看数据并行的上一篇文章，这个过程里产生的单卡通信量是2Φ。
(3) 得到完整的梯度，我们就可以对参数W进行更新。参数W的更新依赖于优化器状态O和梯度。所以每个GPU只能更新对应的参数W。我们需要对参数W进行一次All-Gather，使得所有GPU完成所有参数的更新。产生的单卡通讯量是Φ。
设GPU的个数为N, 我们可以列表去比较一下DDP和优化器优化两者的显存和单卡通讯量，这里我们假设将模型必须存储的数据量为12Φ（K=12），模型参数7.5B，N=64.
在我们的假设下，我们单卡显存降低了4倍，但是单卡通信量只增加1.5倍。看起来是一个不错的改进，那么还可以做的更好吗 ？
Stage 2 梯度划分，每个节点仅保留需要更新自己优化器状态对应的梯度
答案是可以！我们可以将梯度拆分到各个GPU上。
我们刚刚Stage1的分析中，发现由于W的更新依赖于优化器状态和梯度，所以在更新时只能更新相应的参数W。所以Stage 1得到完整的梯度会有一部分的内存冗余，Stage 2对此进行了改进。整体流程如下：
（1）每个GPU保存完整的模型参数W。将一个batch的数据分成n份，分给各个GPU。做完一轮的前向和反向计算后，每个GPU会得到一份梯度。
（2）每个GPU中我们只维护一部分梯度（图4中的绿色部分），所以此时我们只进行Reduce-Scatter操作。目的是使得GPU维护的对应的梯度得到聚合结果。黄色的梯度部分无用，可以丢弃。这个过程的单卡通讯量为Φ。
（3）
每个GPU用自己对应的优化器状态和梯度去更新相应的W。此时对W做一次All-Gather，使得所有GPU完成所有参数的更新。产生的单卡通信量是Φ
图4 梯度划分示意图
