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

Abstract: Large Language Models (LLMs) have revolutionized Natural Language Processing (NLP) but demand massive GPU resources for training. Lowering the threshold for LLMs training would encourage greater participation from researchers, benefiting both academia and society. While existing approaches have focused on parameter-efficient fine-tuning, which tunes or adds a small number of parameters, few have addressed the challenge of tuning the full parameters of LLMs with limited resources. In this work, we propose a new optimizer, LOw-Memory Optimization (LOMO), which fuses the gradient computation and the parameter update in one step to reduce memory usage. By integrating LOMO with existing memory saving techniques, we reduce memory usage to 10.8% compared to the standard approach (DeepSpeed solution). Consequently, our approach enables the full parameter fine-tuning of a 65B model on a single machine with 8 \times RTX 3090, each with 24GB memory. Code and data are available at https://github.com/OpenLMLab/LOMO.

## introduction large language models (llms) have revolutionized natural language processing (nlp), demonstrating remarkable abilities such as emergence and grokking @xcite , pushing model size to become larger and larger. however, training these models with billions of parameters, such as those with 30b to 175b parameters, raises the bar for nlp research. tuning llms often requires expensive gpu resources, such as 8×80gb devices, making it difficult for small labs and companies to participate in this area of research. recently, parameter-efficient fine-tuning methods @xcite , such as lora @xcite and @xcite , provide solutions for tuning llms with limited resources. however, these methods do not offer a practical solution for full parameter finetuning, which has been acknowledged as a more powerful approach than parameter-efficient finetuning @xcite . in this work, we aim to explore techniques for accomplishing full parameter fine-tuning in resource-limited scenarios. we analyze the four aspects of memory usage in llms, namely activation, optimizer states, gradient tensor and parameters, and optimize the training process in three folds: 1) we rethink the functionality of an optimizer from an algorithmic perspective and find that sgd is a good replacement in terms of fine-tuning full parameter for llms. this allows us to remove the entire part of optimizer states since sgd does not store any intermediate state (sec-3.1). 2) our proposed optimizer, lomo as illustrated in figure 1 , reduces the memory usage of gradient tensors to o(1), equivalent to the largest gradient tensor's memory usage (sec-3.2). 3) to stabilize mix-precision training with lomo, we integrate gradient normalization, loss scaling, and transition certain computations to full precision during training (sec-3.3). our technique results in memory usage that equals the usage of parameters plus activation and the largest gradient tensor. we push the memory usage of full parameter fine-tuning to an extreme, making it merely equivalent to the usage of inference. this is because the memory usage of the forward + backward process should not be less than the forward process alone. it is worth noting that, when employing lomo to save memory, we ensure that the fine-tuning process remains uncompromised, as the parameter update process is still equivalent to sgd. we empirically assess the memory and throughput performance of lomo and show that the usage of lomo enables successful training of a 65b model with only 8 rtx 3090 gpus. additionally, to validate the downstream performance of our proposed technique, we apply lomo to tune the full parameters of llms on the superglue dataset collection @xcite . the empirical results demonstrate the efficiency and effectiveness of lomo for optimizing llms with billions of parameters. overall, our contributions are as follows: ## related work in this section, we present related work on memorysaving techniques during full parameter fine-tuning. these techniques can be effectively combined with lomo to further reduce memory consumption. activation checkpointing during vanilla backpropagation, all activations from the forward pass are stored in memory to compute gradients. this can be a significant memory overhead, especially for large language models. alternatively, one could discard all activations and recompute them on demand for gradients computation in order to save memory. however, this can result in a substantial additional computation cost. activation checkpointing (or gradient checkpointing) takes into account both memory usage and computational cost, providing a compromise solution @xcite . the activations of strategically selected checkpoint nodes in the computational graph are kept in memory after the forward pass, while the activations of remaining nodes are recomputed at most once. the activation memory can be reduced to the square root of the original amount at the cost of one extra forward pass. ## experiment in this section, we evaluate our proposed method from three aspects, namely memory profile, throughput and downstream performance. if not further explained, all our experiments are conducted with llama models @xcite , ranging from 7b to 65b. ## conclusion in this paper, we introduce low-memory optimization (lomo), a new optimizer designed to facilitate full parameter fine-tuning for large lan-