Server的通讯量与GPU数量呈线性关系

在数据并行上篇中，我们介绍了朴素数据并行（DP）与分布式数据并行（DDP）。两者的总通讯量虽然相同，但DP存在负载不均的情况，大部分的通讯压力集中在Server上，而Server的通讯量与GPU数量呈线性关系，导致DP一般适用于单机多卡场景。而DDP通过采用Ring-AllReduce这一NCCL操作，使得通讯量均衡分布到每块GPU上，且该通讯量为一固定常量，不受GPU个数影响，因此可实现跨机器的训练。

在上篇介绍中，通讯负载不均的优化我们解释过了，但还遗留了一个显存开销问题：数据并行中，每个GPU上都复制了一份完整模型，当模型变大时，很容易打爆GPU的显存，那要怎么办呢？

今天这篇文章，我们将介绍由微软开发的ZeRO（零冗余优化），它是DeepSpeed这一分布式训练框架的核心，被用来解决大模型训练中的显存开销问题。ZeRO的思想就是用通讯换显存。如果初读ZeRO，觉得它逻辑跳跃，晦涩难懂，那么这篇文章或许可以帮到你～全文结构如下：

一、存储消耗

1.1 存储分类

1.2 混合精度训练

1.3 存储大小

二、ZeRO-DP（），ZeRO与模型并行
三、ZeRO-R
四、ZeRO-offload与ZeRO-Infinity
五、参考

推荐阅读：

图解大模型训练之：流水线并行，以GPipe为例

图解大模型训练之：数据并行上篇（DP, DDP）

一、存储消耗

1.1 存储分类

首先，我们来看在大模型训练的过程中，GPU都需要存什么内容。

存储主要分为两大块：Model States和Residual StatesModel States指和模型本身息息相关的，必须存储的内容，具体包括：

optimizer states：Adam优化算法中的momentum和variance

gradients：模型梯度

parameters：模型参数W

Residual States指并非模型必须的，但在训练过程中会额外产生的内容，具体包括：

activation：激活值。在流水线并行中我们曾详细介绍过。在backward过程中使用链式法则计算梯度时会用到。有了它算梯度会更快，但它不是必须存储的，因为可以通过重新做Forward来算它。

temporary buffers: 临时存储。例如把梯度发送到某块GPU上做加总聚合时产生的存储。

unusable fragment memory：碎片化的存储空间。虽然总存储空间是够的，但是如果取不到连续的存储空间，相关的请求也会被fail掉。对这类空间浪费可以通过内存整理来解决。

1.2 精度混合训练

知道了存储分类，进一步，我们想知道，假设模型的参数W大小是，那么每一类存储具体占了多大的空间呢？

在分析这个问题前，我们需要来了解精度混合训练。

对于模型，我们肯定希望其参数越精准越好，也即我们用fp32（单精度浮点数，存储占4byte）来表示参数W。但是在forward和backward的过程中，fp32的计算开销也是庞大的。那么能否在计算的过程中，引入fp16或bf16（半精度浮点数，存储占2byte），来减轻计算压力呢？于是，混合精度训练就产生了，它的步骤如下图：

存储一份fp32的parameter，momentum和variance（统称model states）

在forward开始之前，额外开辟一块存储空间，将fp32 parameter减半到fp16 parameter。

正常做forward和backward，在此之间产生的activation和gradients，都用fp16进行存储。

用fp16 gradients去更新fp32下的model states。

当模型收敛后，fp32的parameter就是最终的参数输出。

通过这种方式，混合精度训练在计算开销和模型精度上做了权衡。如果不了解fp32，fp16和bf16的细节也没关系，不影响下文的阅读。只要记住它们所占的存储空间和精度表达上的差异即可。

1.3 存储大小

现在，我们可以来计算模型在训练时需要的存储大小了，假设模型的参数W大小是，以byte为单位，存储如下：

因为采用了Adam优化，所以才会出现momentum和variance，当然你也可以选择别的优化办法。因此这里为了更通用些，记模型必存的数据大小为。因此最终内存开销为：

另外，这里暂不将activation纳入统计范围，原因是：

activation不仅与模型参数相关，还与batch size相关

activation的存储不是必须的。存储activation只是为了在用链式法则做backward的过程中，计算梯度更快一些。但你永远可以通过只保留最初的输入X，重新做forward来得到每一层的activation（虽然实际中并不会这么极端）。

因为activation的这种灵活性，纳入它后不方便衡量系统性能随模型增大的真实变动情况。因此在这里不考虑它，在后面会单开一块说明对activation的优化。

二、ZeRO-DP

知道了什么东西会占存储，以及它们占了多大的存储之后，我们就可以来谈如何优化存储了。

注意到，在整个训练中，有很多states并不会每时每刻都用到，举例来说；

Adam优化下的optimizer states只在最终做update时才用到

数据并行中，gradients只在最后做AllReduce和updates时才用到

参数W只在做forward和backward的那一刻才用到

诸如此类

所以，ZeRO想了一个简单粗暴的办法：如果数据算完即废，等需要的时候，我再想办法从个什么地方拿回来，那不就省了一笔存储空间吗？

沿着这个思路，我们逐一来看ZeRO是如何递进做存储优化的。

2.1 : Optimizer State Partitioning

首先，从 optimizer state开始优化。将optimizer state分成若干份，每块GPU上各自维护一份。这样就减少了相当一部分的显存开销。如下图：

复习一下，此时W=fp16，G=fp16，O=fp32。此时，整体数据并行的流程如下：

（1）每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份，做完一轮foward和backward后，各得一份梯度。

（2）对梯度做一次AllReduce，得到完整的梯度G，产生单卡通讯量。为了表达简明，这里通讯量我们就不再换算成byte了，而直接根据参数量来计算。对AllReduce（reduce-scatter + all-gather）不熟悉的朋友，可以先去看上一篇文章。

（3）得到完整梯度G，就可以对W做更新。我们知道W的更新由optimizer states和梯度共同决定。由于每块GPU上只保管部分optimizer states，因此只能将相应的W（蓝色部分）进行更新。（2）和（3）可以用下图表示：

（4）此时，每块GPU上都有部分W没有完成更新（图中白色部分）。所以我们需要对W做一次All-Gather，从别的GPU上把更新好的部分W取回来。产生单卡通讯量。

做完后，设GPU个数为，显存和通讯量的情况如下：

假设各变量大小如表格第二列所示，那么在增加1.5倍单卡通讯开销的基础上，将单卡存储降低了4倍。看起来是个还不错的trade-off，那么还能做得更好吗？

2.2  

现在，更近一步，我们把梯度也拆开，每个GPU格子维护一块梯度。

此时，数据并行的整体流程如下：

（1）每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份，做完一轮foward和backward后，算得一份完整的梯度（下图中绿色+白色）。

（2）对梯度做一次Reduce-Scatter，保证每个GPU上所维持的那块梯度是聚合梯度。例如对GPU1，它负责维护G1，因此其他的GPU只需要把G1对应位置的梯度发给GPU1做加总就可。汇总完毕后，白色块对GPU无用，可以从显存中移除。单卡通讯量。（1）和（2）见下图：

（3）每块GPU用自己对应的O和G去更新相应的W。更新完毕后，每块GPU维持了一块更新完毕的W。同理，对W做一次All-Gather，将别的GPU算好的W同步到自己这来。单卡通讯量。

再次比对下显存和通讯量：

和朴素DP相比，存储降了8倍，单卡通讯量持平，好像更牛皮了呢！那么，还可以优化吗？

2.3  

看到这里，也许你有点感觉了，ZeRO的思想就是：万物皆可切，万物皆可抛。所以现在，我们把参数也切开。每块GPU置维持对应的optimizer states，gradients和parameters（即W）。

数据并行的流程如下：

（1）每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份。

（2）做forward时，对W做一次All-Gather，取回分布在别的GPU上的W，得到一份完整的W，单卡通讯量。forward做完，立刻把不是自己维护的W抛弃。

（3）做backward时，对W做一次All-Gather，取回完整的W，单卡通讯量。backward做完，立刻把不是自己维护的W抛弃。

（4）做完backward，算得一份完整的梯度G，对G做一次Reduce-Scatter，从别的GPU上聚合自己维护的那部分梯度，单卡通讯量。聚合操作结束后，立刻把不是自己维护的G抛弃。

（5）用自己维护的O和G，更新W。由于只维护部分W，因此无需再对W做任何AllReduce操作。

显存和通讯量如下：

到这一步，我们用1.5倍的通讯开销，换回近120倍的显存。只要梯度计算和异步更新做的好，通讯时间大部分可以被计算时间隐藏，因此这样的额外通讯开销，也是划算的。

到这里，我们可以放出原始论文中的说明图了，经过以上分析，这张说明图是不是瞬间就能看懂了。不得不吐槽下，虽然ZeRO的设计不复杂，但对应论文写得真是逻辑跳跃，晦涩难懂...

仔细一想，ZeRO其实掌握了降本增效的精髓：用完即弃，需要再补。反正我补一个和你差不多的，也不会花费很多通（找）讯（人）时间，还大大降低了我的成本。模型的每一层多算（造）几（轮）遍（子）有啥关系呢，反正在我的预算里每个人都一刻不停地干活，就行啦！

2.4 ZeRO VS 模型并行

知道模型并行的朋友，可能会想，既然ZeRO都把参数W给切了，那它应该是个模型并行呀？为什么要归到数据并行里呢？

其实ZeRO是模型并行的形式，数据并行的实质。

模型并行，是指在forward和backward的过程中，我只需要用自己维护的那块W来计算就行。即同样的输入X，每块GPU上各算模型的一部分，最后通过某些方式聚合结果。

但对ZeRO来说，它做forward和backward的时候，是需要把各GPU上维护的W聚合起来的，即本质上还是用完整的W进行计算。它是不同的输入X，完整的参数W，最终再做聚合。

因为下一篇要写模型并行Megatron-LM，因此现在这里罗列一下两者的对比。

三、ZeRO-R

说完了以上对model states的显存优化，现在来看对residual states的优化。

3.1 : Partitioned Activation Checkpointing

前面说过，对activation的存储是灵活的。不像optimizer states，gradients和parameters对模型更新是必须的，activation只是起到加速梯度计算的作用。因此，在哪几层保存activation，保存哪些activation都是可以灵活设置的。同样，我们也可以仿照以上切割方式，每块GPU上只维护部分的activation，需要时再从别的地方聚合过来就行。需要注意的是，activation对显存的占用一般会远高于模型本身，通讯量也是巨大的，所以这块要灵活、有效地实验设计。

3.2 : Constant Size Buffer

固定大小的内存buffer，它的目的在于：

提升带宽利用率。当GPU数量上升，GPU间的通讯次数也上升，每次的通讯量可能下降（但总通讯量不会变）。数据切片小了，就不能很好利用带宽了。所以这个buffer起到了积攒数据的作用：等数据积攒到一定大小，再进行通讯。

使得存储大小可控。在每次通讯前，积攒的存储大小是常量，是已知可控的。更方便使用者对训练中的存储消耗和通讯时间进行预估。

3.3 : Memory Defragmentation

在前文提过，设置机制，对碎片化的存储空间进行重新整合，整出连续的存储空间。防止出现总存储足够，但连续存储不够而引起的存储请求fail。

四、ZeRO-Offload与ZeRO-Infinity

最后，简单介绍一下ZeRO-Offload。它的核心思想是：显存不够，内存来凑。如果我把要存储的大头卸载(offload)到GPU上，而把计算部分放到GPU上，这样比起跨机，是不是能既降显存，也能减少一些通讯压力呢？

ZeRO-Offload的做法是：

forward和backward计算量高，因此和它们相关的部分，例如参数W（fp16），activation，就全放入GPU。

update的部分计算量低，因此和它相关的部分，全部放入CPU中。例如W(fp32)，optimizer states（fp32）和gradients(fp16)等。

具体切分如下图：

ZeRO-infinity也是同理，它们在解决的事情都是：找个除GPU之外的地方，存数据。感兴趣的朋友可以深入研究，这里就不展开了。

　　

      审核编辑：彭静