8-bit最適化アルゴリズム公開した

今月、FaceBookは新しい最適化アルゴリズム公開した、GPUの利用をせっすれると述べる。

• インストール

conda list | grep cudatoolkitかnvcc -V
例えば、君のcudaのバージョンは11.1、
pip install bitsandbytes-cuda111

• 使用

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

リンク：https://github.com/facebookresearch/bitsandbytes