本文介绍了Deepspeed Zero优化器如何解决深度学习模型训练中的显存问题。Zero通过分片策略减少模型参数、梯度和优化器状态的显存占用，分为Zero1、Zero2和Zero3三个档位。Zero1和Zero2的通信量与传统数据并行相当，而Zero3通信量稍大。此外，还提到了Zero-offload技术，用于在GPU显存不足时利用CPU内存。在实践中，Zero1因其良好的显存节省和通信效率平衡而常用。
首先需要明确的是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精度的优化器来进行优化，所以总体其实也没降太多，那么再这种情况，模型的参数占用的显存理论上是优无可优的，因为它其实对应着你实际的卡的数量
最低0.47元/天 解锁文章
816
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
