Hello all,
I recently learned about the SM3 optimization algorithm 29. I implemented it in PyTorch to better understand the paper, and I am sharing the code here. It’s available at PyTorch-SM3 47. The implementation itself includes the features of the TensorFlow version (support for dense and sparse tensors) as well as a feature they mention but do not include in their final version (exponential moving averages). The authors designed the algorithm to be a memory-efficient alternative to Adam and Adagrad for large NLP architectures such as Transformer-Big and BERT.
I have tested my version numerically and tried it for a few simple networks to verify that nothing arose during training. However, I am not working with models as large as the optimizer is intended for. If you do work with models that large and notice any issues, please let me know.