Memory-Efficient Fine-Tuning of Transformers via Token Selection (2024)

Abstract: Fine-tuning provides an effective means to specialize pre-trained models for various downstream tasks. However, fine-tuning often incurs high memory overhead, especially for large transformer-based models, such as LLMs. While existing methods may reduce certain parts of the memory required for fine-tuning, they still require caching all intermediate activations computed in the forward pass to update weights during the backward pass. In this work, we develop TokenTune, a method to reduce memory usage, specifically the memory to store intermediate activations, in the fine-tuning of transformer-based models. During the backward pass, TokenTune approximates the gradient computation by backpropagating through just a subset of input tokens. Thus, with TokenTune, only a subset of intermediate activations are cached during the forward pass. Also, TokenTune can be easily combined with existing methods like LoRA, further reducing the memory cost. We evaluate our approach on pre-trained transformer models with up to billions of parameters, considering the performance on multiple downstream tasks such as text classification and question answering in a few-shot learning setup. Overall, TokenTune achieves performance on par with full fine-tuning or representative memory-efficient fine-tuning methods, while greatly reducing the memory footprint, especially when combined with other methods with complementary memory reduction mechanisms. We hope that our approach will facilitate the fine-tuning of large transformers, in specializing them for specific domains or co-training them with other neural components from a larger system. Our code is available at https://github.com/facebookresearch/tokentune.

## introduction fine-tuning is an effective method for specializing large pre-trained models, either by using direct * equal contribution supervision from the training set of a given task @xcite @xcite , from curated instruction datasets @xcite @xcite , or from human feedback via reinforcement learning @xcite @xcite . however, fine-tuning is not necessarily an efficient method, especially for transformer-based large language models (llms), since their large number of parameters leads to large compute and memory requirements. for instance, fine-tuning gpt-3 175b @xcite or llama 65b @xcite typically requires 1,200 gb and 780 gb of gpu memory, as reported in @xcite and @xcite , respectively. gpu memory usage during fine-tuning can be broken down into three parts: storing (1) the model parameters, (2) the parameter gradients and optimizer states, and (3) the intermediate activations. parameter-efficient fine-tuning (peft) @xcite aims at updating a small number of parameters, e.g., by optimizing a subset of the backbone model's parameters while freezing others, which reduces the memory requirements to store the parameters' gradients and optimizer states. alternatively, quantization techniques @xcite @xcite use low precision data types for model parameters, which reduces the memory cost. for example, in fine-tuning the llama2-7b model, lora @xcite and qlora @xcite , which are representative peft and quantization-based methods, reduce the memory needed for full fine-tuning by 12% and 43%, respectively (figure 1 ). however, such existing approaches still require caching all of the intermediate activations computed in the forward pass to obtain the gradients during the backward pass. in this work, we propose a method for memoryefficient fine-tuning, named tokentune, which aims to significantly reduce the gpu memory dedicated to storing intermediate activations during the forward pass without sacrificing the model performance on various downstream tasks. to this end, tokentune selects a subset of the input tokens in the context, and fine-tunes the model with respect to those selected tokens. more specifically, during the backward pass, tokentune approximates the gradient computation by backpropagating through the selected tokens, and thus only a subset of the intermediate activations need to be cached during the forward pass, thereby reducing the memory cost. we demonstrate the effectiveness of token-tune using both medium-and large-size language models, namely, bert @xcite and llama @xcite , which have hundreds of millions, and billions of parameters, respectively. overall, our results show that fine-tuning with to-kentune leads to downstream task performance on par with that of full fine-tuning or representative methods for memory-efficient fine-tuning, while drastically reducing the memory footprint. notably, tokentune can be effectively combined with existing methods, achieving a greater reduction in memory usage. for instance, by combining to-kentune with qlora @xcite , we can fine-tune llama2-7b using just about one third of the memory qlora alone requires as figure 1 shows. to sum, our contributions are as follows. ## tokentune previous studies analyzing the structure of the sparsity of activations and gradients @xcite @xcite suggest that some neurons and activations could have a predominant importance, while some others may have smaller contributions to the loss and output computation. inspired by these works, we hypothesize that for many downstream tasks, not all tokens in the sequence would need to be involved in the finetuning-more specifically, backpropagation-of transformer models. instead, we conjecture that, when restricted to backpropagating through a subset of tokens, transformers could be further optimized for the downstream task by enabling the additional learning and adjustments, which need to happen during the fine-tuning for the given task, to be done in a more compact way, i.e., by incorporating the additional knowledge more succinctly with respect to the selected subset of tokens. figure 2 illustrates tokentune, aiming at reducing the memory needed to store the intermediate activations used for gradient computation. given an input sequence x, a transformer associates each token from the input sequence to an embedding and computes a corresponding sequence of hidden states h through multiple layer applications. for each input sequence, we select k random positions. 2 we organize each layer's input in two groups, one with the k selected input positions, h g , and the other with the remaining un classification task. the goal is to assign the right class or label y for the given sequence. given the hidden states from the transformer layers, we use the average of the hidden states from the k selected positions of the last layer as input for an mlp, which outputs a probability distribution over the classes of the task, as given by eq. 1. during the evaluation, we use the average from all hidden states of the last layer as input for the mlp. language modeling task. the goal is to learn the probability distribution of a token, given all preceding tokens. we train the language model by applying the traditional cross-entropy loss to the set of k randomly selected positions as given by eq. 2 below, with w lm denoting the head projecting the hidden state back into the vocabulary dimension. the key element of our method is that we disable the gradient computation for the un-selected tokens in ḡ. thus, only the k selected tokens in g contribute to the gradient computation during the backward pass. we detail the method in the case of dense layers and attention mechanism in section 3.1 and section 3.2, respectively. ## application to medium-size encoders alternative methods such as zero-shot learning or prompting usually underperform fine-tuning @xcite . thus, in many cases, finetuning medium size language models may offer a better balance in terms of cost and performance, compared with fine-tuning large language models (llms) or conditioning their outputs with prompt approaches @xcite . medium-size models may also be used as individual components, co-trained to encode information for a larger system @xcite . finally, as detailed in appendix e, the distribution of the gpu memory usage may be very different given the order of magnitude of the fine-tuned algorithm 1: tokentune (we omit layer normalization, skip connections, non-linear functions, and multi-head attention for simplicity) input: input sequence x output: hg, h ḡ 1 compute input token embeddings h 2 re-organize input tokens into two groups (hg and h ḡ ) ## application to large-size decoders we also seek to evaluate our method on larger size pre-trained language models (llms). ## conclusion in this paper, we propose tokentune, a method for reducing the gpu memory required to fine-tune transformer-based models, such as large language models. our contributions are as follows.