optimizer state 都有些啥？_耀世娱乐响应式电子集团

主页 > 新闻中心 > 领导活动

optimizer state 都有些啥？

最近在看显存方面的知识（ZeRO等），很多地方提到optimizer state占用是比较大的一块，例如ZeRO[1]中提到Adam优化器在混合精度训练的时候，占用到12倍的参数大小的显存。

那我有问题了，optimizer states有哪些内容，占多大？

结论：

adam的momentum + variance 占 2*4Byte
混合精度的实现中，需要复制一份fp32的参数作为被optimizer更新的参数， 1*4Byte

总共就是 12倍

先分析一下这两者的大小：

对于SGDM来说，动量和梯度大小一样
对于ADAM来说，多了一个“二阶动量“（variance），和梯度大小一样:

model=torch.nn.BatchNorm2d(10).cuda()
# param初始化
# grad为None
# running_mean,var 初始化

# optimizer=torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
# param_groups: 存在，即weight和bias；
# state_dict()['state']为空，是没有动量的

# print_optimizer_info(optimizer)

input=torch.rand(2,10,10,10).cuda()

output=model(input)
# running_mean,var 更新

optimizer.zero_grad()

output.sum().backward()
# 得到grad
# optimizer.state 还是空

optimizer.step()
# 得到 momentum,
# SGD: optimizer.state为{weight: momentum_buffer, bias: momentum_buffer}
# Adam: optimizer.state为{weight:{exp_avg, exp_avg_sq}, bias:{exp_avg, exp_avg_sq}}

把adam的state打印出来如下：

defaultdict(<class 'dict'>,{Parameter containing:
tensor([0.9000, 0.9000, 1.0998, 1.0999, 0.9002, 0.9000, 0.9001, 0.9001, 0.9001,
        0.9001], device='cuda:0', requires_grad=True): {'step': 1, 'exp_avg': tensor([ 3.8669e-06,  2.3980e-06, -6.4674e-07, -1.0968e-06,  5.2916e-07,
         2.1133e-06,  9.5960e-07,  1.7887e-06,  1.7944e-06,  1.2444e-06],
       device='cuda:0'), 'exp_avg_sq': tensor([1.4953e-12, 5.7502e-13, 4.1827e-14, 1.2030e-13, 2.8001e-14, 4.4662e-13,
        9.2083e-14, 3.1995e-13, 3.2200e-13, 1.5486e-13], device='cuda:0')}, Parameter containing:
tensor([-0.1000, -0.1000, -0.1000, -0.1000, -0.1000, -0.1000, -0.1000, -0.1000,
        -0.1000, -0.1000], device='cuda:0', requires_grad=True): {'step': 1, 'exp_avg': tensor([20., 20., 20., 20., 20., 20., 20., 20., 20., 20.], device='cuda:0'), 'exp_avg_sq': tensor([40., 40., 40., 40., 40., 40., 40., 40., 40., 40.], device='cuda:0')}})

说明momentum 和 var的大小确实和grad一样

optimizer引用weight

optimizer并不会复制一份grad，而是引用。

下面的实验，将model搬到cpu，optimizer的param_group中tensor的device也跟着变了：

(Pdb) p optimizer.param_groups
[{'params': [Parameter containing:
tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], device='cuda:0',
       requires_grad=True), Parameter containing:
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], device='cuda:0',
       requires_grad=True)], 'lr': 0.1, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False}]
(Pdb) model.cpu()
BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(Pdb) p optimizer.param_groups
[{'params': [Parameter containing:
tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True), Parameter containing:
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], requires_grad=True)], 'lr': 0.1, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False}]

上面只是了解了momentum，var的大小，在float32的情况下这俩只占 2*sizeof(float)=8倍参数的大小，还有4倍是什么？

其实答案就在ZeRO文章中，这里截取一部分作为参考：

后面还有具体介绍，就是说，在使用Adam的混合精度训练时，

需要保存以内容：

model本身：

fp16参数一份
fp16梯度一份

optimizer

fp32参数一份
fp32梯度一份（根据实现不同，不一定需要全部保存下来，如果只保留部分梯度，可以忽略）
fp32momentum
fp32variance

这里我其实还有问题：

拷贝fp32的参数是谁做的？真的是optimizer吗？这个在看到Paper 3.1章节之后才理解（后面有解释）
momentum和variance不能都用fp16的吗？

可能是因为精度不够

这里我去查了一下AMP的内容：

混合精度在操作的时候，是先将 FP32 的模型的参数拷贝一份，拷贝的参数转换成 FP16，而 amp 规定了的 FP16 的算子（例如卷积、全连接），对 FP16 的数值进行操作；FP32 的算子（例如涉及 reduction 的算子，BatchNormalize，softmax...），输入和输出是 FP16，计算的精度是 FP32。在反向传播时，依然是混合精度计算，得到数值精度为 FP16 的梯度。最后，由于 GPU 中的 Tensor Core 天然支持 FP16 乘积的结果与 FP32 的累加（Tensor Core math），优化器的操作是利用 FP16 的梯度对 FP32 的参数进行更新

在Mixed precision training | fastai中提到，梯度可以用fp16表示，但是参数更新不能用fp16来做，简单来说，fp16的表示范围有限，会出现 1 + 0.0001=1