A da L omo: Low-memory Optimization with Adaptive Learning Rate (2024)

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 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 and exhibits superior convergence performance compared to LOMO theoretically. To maintain memory efficiency, we employ non-negative matrix factorization for the second-order moment estimation. 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 https://github.com/OpenLMLab/LOMO.

## introduction large language models @xcite @xcite have garnered increasing attention due to their exceptional capabilities across a diverse range of tasks. either supervised fine-tuning or further pre-training can lead to enhanced performance. as the number of parameters grows, the substantial gpu memory required * work done during internship at shanghai ai laboratory. for training sets a high hardware threshold. recently, @xcite has proposed low-memory optimization (lomo) to train large language models in a memory-saving approach by simultaneously backpropagating gradients and updating parameters during the backward pass, enabling the fine-tuning of all parameters of a 7b model on a consumer-grade rtx 3090. while lomo's performance on the super-glue @xcite benchmark is comparable to popular parameter-efficient fine-tuning methods @xcite , it falls short on a broader range of tasks against adaptive optimization methods like adam @xcite , exhibiting a convergence gap. we attribute this to its reliance on the naive stochastic gradient descent optimization approach. we analyze the differences in optimization methods between adam and lomo. compared to lomo, adam incorporates both the first and second moment estimation in its optimizer state, which are the moving averages of the gradient and the squared gradient, respectively. based on our theoretical and empirical analysis, we identify that the second moment estimation is the pivotal factor influencing the convergence of training large language models between lomo and adam. the second-order moment estimation in adam serves to offer an adaptive learning rate for each parameter. expanding on this concept, we introduce the low-memory optimization with adaptive learning rate (adalomo), which similarly provides an adaptive learning rate for each parameter, thus exhibiting superior convergence performance compared to lomo theoretically. to retain memory efficiency, inspired by adafactor @xcite , we employ non-negative matrix factorization @xcite for the second-order moment estimation in the optimizer state. we advocate for the use of a grouped update normalization to stabilize convergence instead of global ## preliminaries in the subsequent sections of this paper, we use θ t to denote the parameters of the model at the t th step of the training process. θ t,i represents the parameter at the i th gradient computation during the backpropagation process of θ t . we use g t to represent the gradient of θ t , and g t,i to denote the gradient of θ t,i . the first and second moment estimation at the t th training step, which are the moving averages of the gradient and the square of the gradient respectively, are represented by m t and v t . the symbol α represents the learning rate. ## method in this section, we introduce our proposed memoryefficient optimization algorithm, adalomo. this algorithm has demonstrated performance comparable to the current de facto optimization method for large language models, adamw, requiring less memory consumption. ## experiments in this section, we evaluate the efficacy of adalomo in instruction-tuning, further pretraining and from-scratch pre-training. additionally, we assess memory usage and throughput. experiments are performed using the llama series of models, which have parameter sizes ranging from 7 billion to 65 billion. ## related work previous research has extensively explored memory-efficient optimizers. adafactor @xcite employs non-negative matrix factorization and approximates the second-order moment estimate @xmath0, achieving sublinear memory consumption. the sm3 algorithm @xcite introduces the cover of the parameters or, more specifically, a set of k non-empty parameter groups. each parameter is assigned an adaptive learning rate based on this cover. for a parameter matrix of size m × n, the sets can be divided by rows and columns, resulting in m + n sets. this reduces the memory requirement from o(m × n) to o(m + n), analogous to adafactor's memory consumption. another line to reduce memory usage is by utilizing low-precision storage for the optimizer state. @xcite and @xcite explored the stability of 16-bit optimizers. the 8-bit optimizer @xcite , using block-wise and dynamic exponent quantization, quantizes the optimizer states of sgdm and adam to 8 bits. the 4-bit optimizer @xcite , employing the newly proposed fp4 format and the adaptive gradient scaling technique. to decrease the memory used by gradients, lomo updates parameters simultaneously during the gradient computation in the backward pass. additionally, there exists a series of memoryefficient optimization methods designed exclusively for fine-tuning. bbt @xcite and bbtv2 @xcite ) utilize evolutionary gradient-free algorithms to optimize continuous prompts without model updates. @xcite employs zeroth-order optimization methods, estimating gradients using two forward passes and optimizing the model in-place, thus equating memory consumption with inference. parameter-efficient fine-tuning (peft) @xcite methods selectively add or pick a subset of parameters for optimization, freezing the majority of the model parameters. in comparison, adalomo updates all parameters using a gradientbased method, suitable for both pre-training and fine-tuning, with memory consumption comparable to peft methods. ## conclusion in this paper, we introduce adalomo, designed to reduce the training barriers for large language models. by incorporating an adaptive learning rate and utilizing grouped update normalization, adalomo achieves results comparable to adamw in instruction-tuning, further pre-training and fromscratch pre-training. concurrently, the memory footprint of adalomo is on par with the peft methods.