S lim F it: Memory-Efficient Fine-Tuning of Transformer-based Models Using Training Dynamics (2024)

Abstract: Transformer-based models, such as BERT and ViT, have achieved state-of-the-art results across different natural language processing (NLP) and computer vision (CV) tasks. However, these models are extremely memory intensive during their fine-tuning process, making them difficult to deploy on GPUs with limited memory resources. To address this issue, we introduce a new tool called SlimFit that reduces the memory requirements of these models by dynamically analyzing their training dynamics and freezing less-contributory layers during fine-tuning. The layers to freeze are chosen using a runtime inter-layer scheduling algorithm. This allows SlimFit to freeze up to 95% of layers and reduce the overall on-device GPU memory usage of transformer-based models such as ViT and BERT by an average of 2.2x, across different NLP and CV benchmarks/datasets such as GLUE, SQuAD 2.0, CIFAR-10, CIFAR-100 and ImageNet with an average degradation of 0.2% in accuracy. For such NLP and CV tasks, SlimFit can reduce up to 3.1x the total on-device memory usage with an accuracy degradation of only up to 0.4%. As a result, while fine-tuning of ViT on ImageNet and BERT on SQuAD 2.0 with a batch size of 128 requires 3 and 2 32GB GPUs, respectively, SlimFit enables fine-tuning them on a single 32GB GPU without any significant accuracy degradation. The code of SlimFit is available at https://github.com/arashardakani/SlimFit.

## introduction over the past few years, various transformer-based models have been developed with the adoption of the attention mechanism that weighs the importance of each part of the input data differently. pre-training of such transformer-based models on large data has led to a significant boost in accuracy when fine-tuned on various natural language processing (nlp) and computer vision (cv) downstream tasks @xcite . despite their great performance in achieving state-of-the-art (sota) accuracy, these models are memory intensive and require a considerably large amount of on-device gpu memory during their fine-tuning phase when compared to the conventional convolutional and recurrent neural networks @xcite . the memory requirement of current transformer-based models has made them difficult to fine-tune even on powerful gpus. with the introduction of larger transformer-based models over the past few years, the on-device gpu memory has become a major bottleneck for their fine-tuning process @xcite @xcite . the total on-device memory usage of gpus consists primarily of activations, parameters, gradients, optimizer states, and the cuda context. among these factors, activations account for most of the memory usage due to batch processing @xcite @xcite as shown in fig. 1 . therefore, activation compressed training (act) has emerged as the primary solution for memory-efficient fine-tuning @xcite . this approach first compresses activations during the forward pass and then decompresses them during the backward pass. in this way, the memory footprint can be significantly reduced by caching the compressed activations. in @xcite @xcite has been a popular choice to compress activations among other compressors such as jpeg @xcite or pruning @xcite . the current sota act adaptively assigns quantization bits to each layer for a given architecture @xcite . while the sota act successfully reduces the memory footprint of activations, its overall on-device gpu memory reduction is not significant. for instance, the total on-device gpu memory reduction of the sota act is limited to 0.1gb despite its 6.4× reduction in the memory of activations when finetuning bert on cola dataset with a batch size of 32. it is worth mentioning that we refer to the memory usage reported by "nvidia-smi" as the overall on-device memory in this paper (see appendix a for more information on memory management). tensor rematerialization @xcite @xcite , also known as gradient checkpointing, is another prominent approach to reducing activation memory by trading computations for memory. in tensor rematerialization, only specific activations are stored during the forward pass, while the rest are recomputed in the backward pass. of course, recomputing activations requires more operations, resulting in a longer fine-tuning process @xcite . reduced precision training, as another approach, performs the computations of both forward and backward passes in low-precision @xcite @xcite . while these works can successfully train conventional models, few-bit model fine-tuning is not trivial. for instance, 8-bit quantization of bert for inference results in a significant precision loss @xcite , which makes fine-tuning on few bits a challenging task. low-rank adaptation (lora) @xcite ) is another key approach to reducing the overall ondevice gpu memory where the transformer-based models are fine-tuned by inserting a small number of trainable parameters into each layer while keeping the pre-trained model parameters frozen. such an approach enables fine-tuning transformer-based models with significantly less number of trainable parameters, leading to a reduction in the memory footprint of optimizer states and gradients. such a memory reduction becomes significant for large transformer models such as gpt @xcite with billions of parameters. different from these methods, we put forward a new approach to reducing the overall on-device memory usage by analyzing training dynamics. more precisely, we dynamically analyze the gradient contributions of layers in transformer-based models and perform parameter updates for specific layers only while the rest of layers are kept frozen. training dynamics have been used to analyze the behavior of a model during its training/fine-tuning process @xcite @xcite . however, our work uses training dynamics to detect and discard unimportant activations during fine-tuning by freezing their associated layers, leading to a reduction of the memory footprint. our method is orthogonal to existing approaches including rematerialization, lora and fused operations @xcite , which could be combined for further reductions. freezing layers or parameters has been studied in different domains, including transformer-based models to preserve previously learned information during fine-tuning @xcite . freezing parameters have also been used to regularize fine-tuning (e.g., over-fitting reduction) in pre-trained models @xcite . recently, freezing has been used to accelerate finetuning by progressively freezing model blocks @xcite @xcite . however, since such an approach starts the fine-tuning process without freezing at least for a few training iterations/epochs, its overall on-device memory requirement remains similar to that of training without freezing. for instance, fine-tuning vit on imagenet with a batch size of 128 using such a freezing approach on a single 32gb gpu results in an out-of-memory error (see appendix b for more details). to orchestrate effective layer-freezing decisions, we introduce a runtime inter-layer scheduling (ils) algorithm. our method finds and freezes a set of layers at each training iteration in transformer-based models that are less contributory, i.e., layers with fewer updates in their parameters, to the fine-tuning process at each iteration. while the ils algorithm successfully detects and freezes unimportant layers, its memory reduction is not proportional to the freezing rate. the reason behind this disproportionality is twofold: the imbalanced number of activations among layers and the existence of static activations. static activations refer to those that cannot be discarded regardless of freezing (e.g., activations of non-linear functions such as gelu). we address these two issues using quantization and pruning to even out the number of activations across all layers and to reduce the memory overhead of static activations. we use quantization and pruning for a few specific layers of transformer-based models as opposed to reduced precision training methods where all the layers are quantized. as a result, the impact of quantization and pruning on accuracy is insignificant in our work. for instance, the accuracy degradation due to quantization and pruning is only 0.1% on the mrpc dataset. by combining ils with quantization and pruning, we introduce a performance tool called slim-fit for reducing the on-device gpu memory usage of transformer-based models during fine-tuning. we demonstrate the effectiveness of slimfit in reducing the memory footprint on popular models of bert and vit. we show that slimfit can freeze up to 95% of layers and reduce the overall on-device memory usage by an average of 2.2× when finetuning bert and vit models on different benchmarks and datasets, such as glue, squad 2.0, cifar-10, cifar-100 and imagenet with an average accuracy degradation of 0.2%. more precisely, slimfit reduces the overall on-device memory usage of the fine-tuning process on glue from 6.1gb to 4.0gb (1.5× reduction) with a batch size of 32, on squad 2.0 from 58.5gb to 19.1gb (3.1× reduction) with a batch size of 128, on cifar-10 from 7.2gb to 4.3gb (1.7× reduction) with a batch size of 32, on cifar-100 from 7.2gb to 4.5gb (1.6× reduction) with a batch size of 32, and on imagenet from 77.4gb to 26.1gb (3.0×) with a batch size of 128 at the cost of up to 0.4% accuracy degradation. as a result, slimfit enables performing memoryintensive fine-tuning processes on a single 32gb gpu such as fine-tuning vit on imagenet with a batch size of 128 while this normally requires three 32gb gpus. ## preliminaries over the past few years, pre-training of attentionbased models has led to significant advances on many nlp and cv tasks with the popular bert @xcite and @xcite models. the pre-training process provides a good initialization point such that these models can better generalize on unseen data of downstream tasks. therefore, these models can achieve stateof-the-art results by fine-tuning through small adjustments to their parameters. architecturally, these models consist of an initial embedding layer, followed by repeated blocks of multi-head attention (mha) fed into a feed-forward network (ffn) module (see appendix c for more details). the base architectures of bert and vit contain over a hundred layers built up in this manner. despite the large number of layers, not all need to be updated during fine-tuning to achieve decent performance on downstream tasks, as shown in @xcite . notably, the authors found that freezing approximately 60% of early attention layers in bert led to negligible performance degradation. this suggests that the fine-tuned model tends to preserve generic features learned during pre-training. motivated by this study, we seek to analyze the training dynamics of pre-trained models and to automatically detect layers with less contributions to the fine-tuning process. ## learning the importance of layers training dynamics is an active field of research that provides insight about the behavior of pre-trained models when fine-tuning on downstream tasks. the convergence proof of optimization algorithms such as stochastic gradient descent @xcite shows that the distance between the parameters and the optimal solution is reduced over training iterations and accordingly, the weight distance (or the weight update amount) between consecutive iterations decreases. therefore, it is possible that some layers can only receive minimal changes to their parameters as we approach the end of the training process. of course, detecting and freezing such layers, when they show minimal updates, will not affect accuracy. since transformerbased models are pre-trained, they already show small updates during fine-tuning compared to pretraining. as such, detecting and freezing layers with minimal updates (i.e., weight distance values) will not significantly affect the fine-tuning process and accordingly the final accuracy. based on the above observations, we consider the ℓ 1 -norm of the update received by parameters of each layer through all the fine-tuning iterations as the training dynamics in this paper. it is also worth mentioning that freezing layers has no impact on training convergence as it causes a pause in the training procedure of frozen layers as shown by our theoretical analysis in appendix d.1. ## inter-layer load-balancing so far, we have introduced our ils algorithm that prioritizes updating particular layers while keeping the rest of layers frozen according to their distance value. for the given freezing rate of 50% as an example, we expect to see a 2× reduction in the memory footprint of activations. however, this is not the case in transformer-based models due to the imbalanced the number of activations across all the layers. in fact, the imbalance in the number of activations undermines the ability of our ils algorithm in reducing the memory footprint during the fine-tuning as shown in fig. 4b . since the focus of this paper is on transformerbased models such as bert and vit, we analyze their architecture for imbalanced layers. ## dynamic and static activations assuming that the backpropagation is performed from the last to the first layer, the type of activations in transformer-based models can be divided into two categories: dynamic and static. we refer to the activations that can be discarded by freezing their layer as dynamic activations. on the other hand, static activations cannot be discarded regardless of freezing. among different types of layers, gelu, matmul, softmax and layernorm contain static activations as shown gelu and matmul/softmax do not have any trainable parameters and accordingly cannot be frozen. therefore, these two layers hold on to their activations throughout the fine-tuning process. the best approach to reduce their memory cost is quantization. we use 4 and 8 bits for quantization of activations in gelu and matmul/softmax, respectively. since there is no 4-bit tensor support in pytorch, we store each two 4-bit activations as a single 8-bit activations using shift operations. note that using such bit-levels result in a negligible accuracy degradation while further quantization of those activations incurs a significant accuracy loss. as opposed to gelu and matmul/softmax, lay-ernorm contains trainable parameters and can be frozen by the ils algorithm. however, its activations are still static. the forward pass of layernorm is computed by: ## slimfit slimfit is a performance tool that exploits our ils algorithm along with quantization and pruning to reduce the memory footprint of activations through an iterative freezing process. the total on-device gpu memory reduction of slimfit is a result of the memory reduction in both dynamic and static activations. static activations contribute a fixed amount of memory whereas the memory usage of dynamic activations depends on the freezing rate. given a high freezing rate, the memory footprint of activations and accordingly the total on-device gpu memory usage can be significantly reduced. the choice of freezing rate depends on the memory budget of 6223 the user. by increasing the freezing rate up to a certain point, there will be no performance degradation. however, using an extremely high freezing rate trades off memory for accuracy. finding the breaking point of the method is task dependent and varies from one dataset to another. ## experimental results we use the base version of bert and vit for our experiments. we fine-tune these two models using slimfit which is implemented on pytorch. we evaluate bert @xcite using the glue benchmark @xcite and squad 2.0 @xcite . for vit (dosovitskiy et al., 2021), we use cifar-10, cifar-100 and imagenet datasets @xcite for evaluation purposes. we discuss the memory usage of activations and the overall on-device gpu memory on the 32gb nvidia v100 gpu. we report the total on-device gpu memory usage using "nvidia-smi". for all the experiments in this section, we use 3 epochs for fine-tuning. the details about the cv/nlp tasks, measurements and hyper-parameter settings are provided in appendix g. ## ablation studies due to limited space, we only discuss the impact of the freezing rate on accuracy of slimfit in this section. we provide further discussions on the impact of quantization/pruning, the total wall-clock time and the frequency of update occurrence in appendices i, j, and k, respectively. for all the experiments in this section, we use a batch size of 32 and 3 epochs for fine-tuning. our ils algorithm orchestrates the freezing schedule based on a simple rule: layers with largest distance values are updated whereas those with lowest distance values are kept frozen for the given freezing rate. of course, such an iterative freezing approach trades off between accuracy and freezing rate. to better show this trade-off, we measured and illustrated accuracy of cola and mrpc datasets across different freezing rates in fig. 6 . the trade-off curve shows our ils algorithm can maintain the accuracy at the same level of the baseline by freezing up to 95% of layers. besides our ils algorithm, the freezing schedule can be decided using random or progressive freezing approaches. in the random scheduling method, frozen layers are randomly selected at each iteration. in the progressive approach, on the other hand, early layers are progressively kept frozen whereas later layers are being updated throughout the finetuning process. among these approaches, our ils algorithm significantly stands out in terms of both accuracy and freezing rate as shown in fig. 6 . the reason behind its superior performance is that ils allows more updates for layers with large distance values by keeping layers with minimal distance values frozen for a specific number of iterations. on the other hand, in the random approach, the layers are randomly selected to be updated. therefore, layers with large distance values receive less number of updates in the random approach compared to ils. of course, the chance of layers with large distance values being randomly selected as active layers decreases as the freezing rate increases, which explains the accuracy gap between ils and the random approach with freezing rate higher than 70% freezing rate. in the progressive freezing approach, the early layers receive no update during the fine-tuning process, resulting in a significant accuracy degradation for large freezing rates. ## comparison with sota techniques next, we compare slimfit with state-of-the-art compression methods targeting memory reduction, i.e., 4-bit gact @xcite and @xcite . ## conclusion in this paper, we presented a performance tool called slimfit that reduces the memory usage of activations and accordingly the overall on-device gpu memory usage of transformer-based models through an iterative freezing of layers during fine-tuning. slimfit adopts an inter-layer scheduling method to orchestrate the freezing schedule at each iteration. to balance the number of activations across all layers and to reduce the memory usage of static activations, slimfit uses quantization and pruning for a few specific layers. we evaluated the performance of slimfit across different nlp and cv tasks. we showed that slimfit significantly reduces the on-device gpu memory usage of the finetuning process by up to 3.1× when using a batch size of 128.