Pushing the Limits of Low-Bit Optimizers: A Focus on EMA Dynamics
This addresses the problem of prohibitively high training costs for large-scale models by reducing optimizer memory usage, though it is incremental as it builds on existing Adam-style optimizers.
The paper tackles the high memory cost of stateful optimizers like Adam by introducing SOLO, a novel optimizer that reduces memory footprint by enabling quantized states with as low as 2-3 bits, achieving substantial memory savings with minimal accuracy loss.
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.