Facebook 推出 8 比特优化器，可以节省75%的显存

最近，Facebook 推出了支持 pytorch 的 8 位优化器，在减小内存占用的同时，竟然还能保持和32位优化器相当的准确性。

facebook yyds！！！！

论文链接:

https://arxiv-download.xixiaoyao.cn/pdf/2110.02861.pdf

开源链接:

https://github.com/facebookresearch/bitsandbytes

安装

硬件要求：NVIDIA Maxwell GPU 或更新版本 (>=GTX 9XX) 支持的 CUDA 版本：9.2 - 11.3

# choices: {cuda92, cuda 100, cuda101, cuda102, cuda110, cuda111, cuda113}
# replace XXX with the respective number
pip install bitsandbytes-cudaXXX

很简单吧！ 等着安装完成。。。。

使用

使用也很简单，直接用bitsandbytes替换原来的优化器即可。

import bitsandbytes as bnb
# adam = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.995)) # comment out old optimizer
adam = bnb.optim.Adam8bit(model.parameters(), lr=0.001, betas=(0.9, 0.995)) # add bnb optimizer
adam = bnb.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.995), optim_bits=8) # equivalent
torch.nn.Embedding(...) ->  bnb.nn.StableEmbedding(...) # recommended for NLP models

总结

关于优化器的论文还是比较少的，目前主流的还是Adam和SGD。今天看到这个我测试了一下。

首先用win10 的电脑测试，发现不能用。

报错了，应该是不支持，.so是linux的文件，所以更换到Ubuntu16.04上尝试，环境：Ubuntu16.04、CUDA10.2。

安装命令：

pip install bitsandbytes-cuda102

1

我先使用AdamW尝试，模型是TextRNN，16个BatchSize显存占用是9G左右，测试了1个epoch。

然后在使用Adam8bit，16个BatchSize显存占用也是9G左右，并没有减少。

有点遗憾！本想着能解决显卡不够的问题，看来还是不行！

唉！！！！！！！

也需是使用的姿势不对吧！希望有人能做更详细的解答！