Pushing the Limits of Low-Bit Optimizers: A Focus on EMA Dynamics

TL;DR

This paper introduces SOLO, a low-memory optimizer that quantizes Adam-style states to as low as 2-3 bits, significantly reducing memory use while maintaining training effectiveness.

Contribution

SOLO is a novel optimizer enabling ultra-low-bit quantization of Adam states, addressing key challenges with tailored solutions for improved efficiency.

Findings

Achieves 2-3 bit quantization of Adam states.

Maintains training accuracy with minimal loss.

Substantial memory savings demonstrated.

Abstract

The rapid scaling of models has led to prohibitively high training and fine-tuning costs. A major factor accounting for memory consumption is the widespread use of stateful optimizers (e.g., Adam), which maintain auxiliary information of even 2x the model size in order to achieve optimal convergence. We therefore present SOLO in this work to spawn a novel type of optimizer that requires an extremely light memory footprint. While previous efforts have achieved certain success in 8-bit or 4-bit cases, SOLO enables Adam-style optimizers to maintain quantized states with precision as low as 3 bits, or even 2 bits. This immense progress is due to the identification and resolution of two key challenges: the signal swamping problem in unsigned quantization that results in unchanged state dynamics, and the increased gradient variance in signed quantization that leads to incorrect descent directions. The theoretical analysis suggests a tailored logarithmic quantization for the former and a precision-specific momentum hyperparameter for the latter. SOLO can thus be seamlessly applied to Adam-style optimizers, leading to substantial memory savings with minimal accuracy loss.