ICML 2024 | Gradient checkpointing too slow? Without slowing down and saving video memory, LowMemoryBP greatly improves backpropagation video memory efficiency

The AIxiv column is a column where this site publishes academic and technical content. In the past few years, the AIxiv column of this site has received more than 2,000 reports, covering top laboratories from major universities and companies around the world, effectively promoting academic exchanges and dissemination. If you have excellent work that you want to share, please feel free to contribute or contact us for reporting. Submission email: liyazhou@jiqizhixin.com; zhaoyunfeng@jiqizhixin.com

The first author of this paper is Yang Yuchen, a second-year master's student in the School of Statistics and Data Science of Nankai University, and his advisor is Associate Professor Xu Jun in the School of Statistics and Data Science of Nankai University. The research focus of Professor Xu Jun’s team is computer vision, generative AI and efficient machine learning, and has published many papers in top conferences and journals, with more than 4,700 Google Scholar citations.

Since large-scale Transformer models have gradually become a unified architecture in various fields, fine-tuning has become an important means of applying pre-trained large models to downstream tasks. However, as the size of the model increases day by day, the video memory required for fine-tuning also gradually increases. How to efficiently reduce the fine-tuning video memory has become an important issue. Previously, when fine-tuning the Transformer model, in order to save graphics memory overhead, the usual approach was to use gradient checkpointing (also called activation recalculation) to reduce the time required in the backpropagation (BP) process at the expense of training speed. Activate video memory usage.

Recently, the paper "Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation" published at ICML 2024 by the team of Teacher Xu Jun from the School of Statistics and Data Science of Nankai University proposed that by changing the backpropagation (BP) process, in Without increasing the amount of calculation, the peak activation memory usage is significantly reduced.

• Paper: Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation

• Paper link: https://arxiv.org/abs/2406.16282

• Project link: https:// /github.com/yyyyychen/LowMemoryBP

The article proposes two backpropagation improvement strategies, namely Approximate Backpropagation (Approx-BP) and Memory-Sharing Backpropagation (MS-BP). Approx-BP and MS-BP respectively represent two solutions to improve memory efficiency in backpropagation, which can be collectively referred to as LowMemoryBP. Whether in a theoretical or practical sense, the article provides groundbreaking guidance for more efficient backpropagation training.

In theoretical memory analysis, LowMemoryBP can significantly reduce the activation memory usage from activation functions and normalization layers. Taking ViT and LLaMA as examples, fine-tuning ViT can reduce activation memory by 39.47%, and fine-tuning LLaMA can reduce activation by 29.19%. video memory.

In actual experiments, LowMemoryBP can effectively reduce the peak memory usage of Transformer model fine-tuning including ViT, LLaMA, RoBERTa, BERT, and Swin by 20%~30%, and will not increase the training throughput and Loss of test accuracy.

Approx-BP

In traditional backpropagation training, the backpropagation of the gradient of the activation function strictly corresponds to its derivative function. For the GELU and SiLU functions commonly used in the Transformer model, this means that the input needs to be The feature tensor is completely stored in the active video memory. The author of this article proposed a set of backpropagation approximation theory, namely Approx-BP theory. Guided by this theory, the author uses a piecewise linear function to approximate the activation function, and replaces the backpropagation of the GELU/SiLU gradient with the derivative of the piecewise linear function (step function). This approach leads to two asymmetric memory-efficient activation functions: ReGELU2 and ReSiLU2. This type of activation function uses a 4-stage step function for reverse passback, so that the activation storage only needs to use a 2-bit data type.

MS-BP

BP Each layer of the network usually stores the input tensor into the activation memory for backpropagation calculation. The author points out that if the backpropagation of a certain layer can be rewritten into an output-dependent form, then this layer and the subsequent layer can share the same activation tensor, thus reducing the redundancy of activation storage.

The article points out that LayerNorm and RMSNorm, which are commonly used in Transformer models, can well meet the requirements of the MS-BP strategy after merging affine parameters into the linear layer of the latter layer. The redesigned MS-LayerNorm and MS-RMSNorm no longer generate independent active graphics memory.

Experimental results

The author conducted fine-tuning experiments on several representative models in the fields of computer vision and natural language processing. Among them, in the fine-tuning experiments of ViT, LLaMA and RoBERTa, the method proposed in the article reduced the peak memory usage by 27%, 29% and 21% respectively, without causing any loss in training effect and training speed. Note that the comparison Mesa (an 8-bit Activation Compressed Training method) reduces the training speed by about 20%, while the LowMemoryBP method proposed in the article completely maintains the training speed.

Conclusion and significance

The two BP improvement strategies proposed in the article, Approx-BP and MS-BP, both achieve the activation of video memory while maintaining the training effect and training speed. significant savings. This means that optimizing based on the BP principle is a very promising memory saving solution. In addition, the Approx-BP theory proposed in the article breaks through the optimization framework of traditional neural networks and provides theoretical feasibility for using unpaired derivatives. Its derived ReGELU2 and ReSiLU2 demonstrate the important practical value of this approach.

You are welcome to read the paper or code to understand the details of the algorithm. The relevant modules have been open sourced on the github repository of the LowMemoryBP project.

The above is the detailed content of ICML 2024 | Gradient checkpointing too slow? Without slowing down and saving video memory, LowMemoryBP greatly improves backpropagation video memory efficiency. For more information, please follow other related articles on the PHP Chinese website!