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

Posted 2021-12-14 AI浩

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

我先使用AdamW尝试，模型是TextRNN，16个BatchSize显存占用是9G左右，测试了1个epoch。
然后在使用Adam8bit，16个BatchSize显存占用也是9G左右，并没有减少。
有点遗憾！本想着能解决显卡不够的问题，看来还是不行！
唉！！！！！！！
也需是使用的姿势不对吧！希望有人能做更详细的解答！
（图片在服务器上，明天补上图！）