AdaLomo: Low-memory Optimization with Adaptive Learning Rate

Kai Lv,Hang Yan,Qipeng Guo,Haijun Lv,Xipeng Qiu

2024-06-06

Abstract:Large language models have achieved remarkable success, but their extensive parameter size necessitates substantial memory for training, thereby setting a high threshold. While the recently proposed low-memory optimization (LOMO) reduces memory footprint, its optimization technique, akin to stochastic gradient descent, is sensitive to hyper-parameters and exhibits suboptimal convergence, failing to match the performance of the prevailing optimizer for large language models, AdamW. Through empirical analysis of the Adam optimizer, we found that, compared to momentum, the adaptive learning rate is more critical for bridging the gap. Building on this insight, we introduce the low-memory optimization with adaptive learning rate (AdaLomo), which offers an adaptive learning rate for each parameter. To maintain memory efficiency, we employ non-negative matrix factorization for the second-order moment estimation in the optimizer state. Additionally, we suggest the use of a grouped update normalization to stabilize convergence. Our experiments with instruction-tuning and further pre-training demonstrate that AdaLomo achieves results on par with AdamW, while significantly reducing memory requirements, thereby lowering the hardware barrier to training large language models. The code is accessible at <a class="link-external link-https" href="https://github.com/OpenLMLab/LOMO" rel="external noopener nofollow">this https URL</a>.

Machine Learning,Computation and Language

What problem does this paper attempt to address?

The paper mainly discusses how to reduce memory consumption when training large language models. Although existing Low Memory Optimizations (LOMO) methods have reduced memory footprint, their optimization effects still lag behind the popular AdamW optimizer. It is found that the adaptive learning rate in Adam optimizer is crucial for the convergence of large language models. Therefore, the paper proposes AdaLomo, a low memory optimization method that combines adaptive learning rate. AdaLomo provides an adaptive learning rate for each parameter, theoretically having better convergence performance than LOMO. To maintain memory efficiency, AdaLomo adopts non-negative matrix decomposition to estimate the second moment and uses group-wise updated normalization to stabilize the training process, which is faster than global updated normalization and does not affect performance. Experimental results show that AdaLomo performs comparably to AdamW in fine-tuning, further pre-training, and from-scratch pre-training, but significantly reduces memory requirements, lowering the hardware threshold for training large language models. The paper also analyzes the comparison between AdaLomo and other optimizers in terms of memory consumption and throughput, demonstrating the effectiveness and efficiency of AdaLomo.

AdaLomo: Low-memory Optimization with Adaptive Learning Rate

Full Parameter Fine-tuning for Large Language Models with Limited Resources

APOLLO: SGD-like Memory, AdamW-level Performance

BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models

CAME: Confidence-guided Adaptive Memory Efficient Optimization

AdaRankGrad: Adaptive Gradient-Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning

LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics

Adaptive Optimization for Enhanced Efficiency in Large-Scale Language Model Training

Promoting Exploration in Memory-Augmented Adam using Critical Momenta

Efficient Adaptive Optimization via Subset-Norm and Subspace-Momentum: Fast, Memory-Reduced Training with Convergence Guarantees

ZO-AdaMU Optimizer: Adapting Perturbation by the Momentum and Uncertainty in Zeroth-Order Optimization

AdaZeta: Adaptive Zeroth-Order Tensor-Train Adaption for Memory-Efficient Large Language Models Fine-Tuning

Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark

BlockLLM: Memory-Efficient Adaptation of LLMs by Selecting and Optimizing the Right Coordinate Blocks

Addax: Utilizing Zeroth-Order Gradients to Improve Memory Efficiency and Performance of SGD for Fine-Tuning Language Models

LoMA: Lossless Compressed Memory Attention

OwLore: Outlier-weighed Layerwise Sampled Low-Rank Projection for Memory-Efficient LLM Fine-tuning

Adam-mini: Use Fewer Learning Rates To Gain More

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training

Dynamic Memory Based Adaptive Optimization

GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection