[论文阅读] ZeRo

zero主要用于数据并行

对于大模型来说,内存主要用来保存模型状态,主要有

  1. 模型参数,梯度(梯度和模型大小是一样的),优化器状态(adam中需要保存momentum和variance信息)。
  2. 所有中间值(gpipe使用重算解决这个问题),临时buffer,内存碎片化

半精度训练:
模型某层的w,和层的输入输出,都使用fp16
例如,Wx = y ,这三者都是fp16

但是有一个问题,fp16通常精度不够,加着加着就炸了,或者,太小的情况一直加都是0。
这个情况在累积权重(不断把梯度的东西加进来)的时候会出现,所以如果权重也是fp16的话,可能加半天都加不动。

所以权重更新使用fp32,也即是权重还有一个fp32的copy,作梯度更新的时候使用fp32,然后转成fp16,再参与前向后向计算。

假设模型n个可学习参数,那存fp16的参数和梯度各需要2n和2n个bytes,(一个fp16需要2个byte,1个byte是8bit),还有fp32的内存是权重,momentum,variance,各需要4n,4n,4n,共12n。

所以计算前向和反向只需要4n个bytes,但是为了更新的时候需要存12n个bytes。

zero 1

12ψ / n,n为多少张卡

上一篇:spring boot sso-http流程


下一篇:深入剖析C++多态的实现与原理-详解