zero

对于模型的训练，训练的速度和显存的占用是必须要考虑的两个因素，特别是现在模型越来越大。1.4B的模型，在32GB的GPU上训练就会OOM。更别提现在动不动就几百B甚至上千B的模型。所以分析那些因素对模型的训练速度和显存的占用是十分必要的。 显存占用分析（训练阶段） 在训练阶段，显存被如下组件占用 model weights optimizer states gradients forward activations saved for gradient computation temporary buffers functionality-specific memory 在ZeRO中model weights、optimizer states、gradients被称为模型状态（model states）, 剩下的被称为剩余状态（residual states） 具体的计算如下（参数量假设为1） model weights 4 bytes ： fp32 training 6 bytes ： mixed precision training（即需要保存一个float32参数，又需要保存一个float16参数） Optimizer States 8 bytes：对于大模型优化器一般为AdamW（包含一阶梯度和二阶梯度，所以对于一个参数，优化器占用8个比特） 2 bytes：8-bit AdamW optimizer 4 bytes：SGD with momentum Gradients 4 bytes： fp32 or mixed precision training （注：对于混合精度训练，一个参数的梯度，ZeRO论文任务是2 bytes(float16), Hugging face中认为梯度一般是4 bytes(float32)。）。所以这里不太确定，获取两种计算方式都是正确的（由框架实现决定） 所以，如果使用混合精度训练，一个参数，需要消耗18个bytes（6+8+4）（ZeRO认为16个bytes） 减少显存使用和提升训练速度的tricks Method Speed Memory 备注 Gradient accumulation No Yes Gradient checkpointing No Yes Mixed precision training Yes (No) 不太严谨 Batch size Yes Yes Optimizer choice Yes Yes DataLoader Yes No DeepSpeed Zero No Yes 必要的解释...