首先需要明确的是Deepspeed本身是一种数据并行的优化，它也可以和其他的PP，TP，一起结合使用。

       Deepspeed最有名的feature就是大名鼎鼎的Zero，我们之前讲过，在训练的过程中占用显存的数据主要分两类：

      一类是模型本身的占用显存，如果用一个正常的混合精度训练的话，那么需要16byte，也就是2字节的模型参数，2字节的模型梯度，如果是以Adam来做优化器的话，那么要以32byte分别存取，Adam的状态，Adam的变量momentum和变量variance，这些一共耗费12字节，也就是一个模型的参数要消耗掉16字节的显存存储空间。

      第二类是Activation，其实严格说应该叫residual status，包含激活值Activation，各种临时的buffer，还有无法使用的碎片 fragments，这里面最大头就是Activation。

     如果Activation撑不住了，实在不行可以采用Activation checkpoint的方式，让钱箱传播的时候所有的激活都存，在反向传播的时候重算一次，也就是通过时间来空间，这个虽然导致整个训练过程会变慢，但是起码还可以正常进行。

     另外的第一类就实在是没办法缩减了，怎么缩减呢，有两种方式，第一是用FP8来训练，但是这个取决于你的卡支持的算子，目前看就H100支持FP的训练，另外就是即使用FP8.我们也会使用FP32精度的优化器来进行优化，所以总体其实也没降太多，那么再这种情况，模型的参数占用的显存理论上是优无可优的，因为它其实对应着你实际的卡的数量。

      Deepspeed Zero (Zero Redundancy Optimizer) 主要针对的就是第一类占用显存的优化，它设计的目的就是为了完成这部分优化的。

     首先看一下Deepspeed的设计理念，主要还是分片，在这个角度上它和标准的模型并行的理解并无二致，但是比如一个70B的模型，以BF16/FP16来进行训练，这个是根本不可能开启DDP的，因为单卡的显存消耗就达到了140G，不算Activation。

      这个时候当开启了Deepspeed Zero，整个的显存占用就不一样了。

Zero

如图 所示，Zero分成若干个档:

那么第一档叫做Zero1 Pos，os指的是optimize status, 就是，看刚才列的式子，最应该优化的参数是哪一个? 肯定首当其冲的就是Adam这个优化器。还是以70B的模型举例。此时，模型参数和梯度，这两部分，仍旧是每张卡保持一份，每张卡是35G，但是Adam则是被分成了2份，每张卡就存一份，也就是12byte/N, N=卡数，这个例子中N=2，原来每张卡要存的这部分ADAM是105/2 是52.5, 这样一张卡要承载87.5G的显存，显然是A100 80G的显卡是不够的，所以N这个值就非常关键了，在Zero中N越大，单张显卡节省的内存就越多，这两者是正比关系，因为被拆分的Adam是分子，N是分母所以当N远远大于12的时候，就越来越趋近于单卡上只剩下了模型参数和梯度，趋向于每参数只有4字节的占用，也就是原来16字节的1/4。

第二档叫做Zero2 Pos+g，g就是gradient就是梯度的意思，顾名思义，梯度也被加入到分拆的行列，而N远大于14的时候，单卡上就只剩下了没参数2字节的占用。也就是原来的1/8。

第三档就是Zero3 Pos+g+p ，p 就是Parameter，所以Zero3连模型参数也给拆分了，这个如果N特别大，可以认为显存无限趋近于0了。

看起来Zero几乎是面临显存限制时的最好解法，在不考虑通信的前提下。

现在需要把通信的代偿计算进来，可以先说明结论，Zero1和Zero2与传统的数据并行所占用的通信量是一至的。

传统的DDP这样的数据并行，在每一步计算梯度以后，需要通过All-reduce来计算梯度的均值，分为ReduceScatter加上AllGather这两部分吗，这个时候每张卡在发送和接收两个方向上的通信量是2*模型参数。

因为Zero2会拆分更细粒度，所以直接以Zero2为例来讲解，在Zero2的环境下，每张卡只存储了1/N的优化器和梯度，对于本地的GPU卡来讲，为了计算这1/N的梯度的均值，需要进行一次Reduce操作，通信量是1/N*parameter*N, 也就是参数量，因为其他的显卡也不需要保留这部分梯度值，所以每张卡就只需要发送一次即可。

计算好梯度以后，就要更新本地的优化器了，反向传播伴随着一次Gather的操作，通信量和上面一样也等于参数量。

两者加起来就是2*模型参数。

而Zero3就不一样了，因为每张卡只存了1/N的参数，所以涉及到参数的同步，也就多了一次broadcast的操作，所以它的通信量就是3*模型参数量，也就是DDP，Zero1, Zero3的1.5倍。

Zero-offload，也有被称为Zero4的, 这个甚至是一张卡都可以，Zero-offload就是针对GPU显存不够，甚至单卡的场景设计出来的。

我们看一个混合精度训练的场景：

图  混合精度训练-1

比如图是某一层的训练iteration，在前向计算的时候，要用到上一层的activation，和本层的参数，反向传播求导的时候，也要用到相同的东西来求梯度，同时Adam优化器对权重参数进行更新，假设模型参数为M，那么在混合精度下进行训练，要么是4M的参数参与要么是2M的参数参与计算。

Zero-offload是采用CPU的内存来顶替GPU显存不足的一种方式，在考虑那些计算的步骤放在CPU上的时候，我们先看一下计算复杂度，图中有几个计算节点：前向计算FWD，反向传播BWD，参数更新 param update，优化器更新权重float2half。

这其中FWD和BWD的计算复杂度要高，可以认为是对模型参数M的操作乘上batchsize，而后两个的计算复杂度实际上就是等于对模型参数的操作。

图 混合精度训练-2

如图所示，Zero-offload，为了优化各种算子的执行，把复杂的FWD和BWD放在了GPU上执行，而相对固定可控的算子param update和float2half就放在了CPU上进行运算，而这一部分除了相对算子简单以外，最重要的是Adam优化器本身是32bit的，把它放在CPU的内存上，会极大的节省GPU的显存。

在多卡的Zero-offload的情况下，采用了Zero2的方式进一步节省GPU的内存。

在现实项目里，更多用的实际上是Zero1，因为Zero1已经省去了3/4的显存开销，而且需要All-Reduce同步的信息又相对合理，是节省显存和节省通信的最好的方式。

     除了Zero以外，Pytorch 本身有自己的FSDP，来实现类似的功能，FSDP可以看成是ZERO-3的实现，传统的数据并行(DDP)是在每一个GPU卡上保存整个model的参数/梯度/优化器状态, 然后对数据集切分为 N 个shard分片给不同的GPU进行训练，计算完梯度后通过all-reduce通信来做梯度的融合。

在FSDP中的主要思路是想办法把model的梯度/优化器状态/参数都进行切分操作，每个GPU只存最少得信息，1/N，也就是在ZERO-3的思路。核心是把DDP中的all-reduce操作拆解为reduce-scatter和all-gather 操作。

FSDP在Pytorch 1.11之后的版本可以使用。