Reproduce “Memory Efficient Optimizers with 4-bit States” on OpenResearch

About this paper

Bingrui Li, Jianfei Chen, Jun Zhu

Researchers at Tsinghua University developed 4-bit precision optimizers for neural network training, achieving up to 57.7% memory savings for large language models while maintaining performance on par with 32-bit counterparts. This enables training larger models or batch sizes on existing hardware by tackling specific challenges in low-bit quantization.

Key insights

• Gradient first moments exhibit heterogeneous outlier patterns across different network architectures and layers, necessitating finer-grained quantization than simple block-wise approaches.
• Quantizing the second moment requires specialized handling due to a critical "zero-point problem" where mapping small non-zero values to zero can destabilize the Adam update rule.
• Different moment statistics (first and second) possess distinct distributional properties and outlier characteristics, demanding tailored quantization methodologies rather than a one-size-fits-all approach.

Summary by alphaXiv.