首先需要明确的是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精度的优化器来进行优化,所以总体其实也没降太多,那么再这种情况,模型的参数占用的显存理论上是优无可优的,因为它其实对应着你实际的卡的数量