Mesa: A Memory-saving Training Framework for Transformers

Abstract

There has been an explosion of interest in designing high-performance Transformers. While Transformers have delivered significant performance improvements, training such networks is extremely memory intensive owing to storing all intermediate activations that are needed for gradient computation during backpropagation, especially for long sequences. To this end, we present Mesa, a memorysaving resource-efficient training framework for Transformers. Specifically, Mesa uses exact activations during forward pass while storing a low-precision version of activations to reduce memory consumption during training. The low-precision activations are then dequantized during backpropagation to compute gradients. Besides, to address the heterogeneous activation distributions in the multi-head self-attention layers, we propose a head-wise activation quantization strategy, which quantizes activations based on the statistics of each head to minimize the approximation error. To further boost training efficiency, we learn quantization parameters by running estimates. More importantly, by re-investing the saved memory in employing a larger batch size or scaling up model size, we may further improve the performance under constrained computational resources. Extensive experiments on ImageNet, CIFAR-100 and ADE20K demonstrate that Mesa can reduce half of the memory footprints during training while achieving comparable or even better performance. Code is available at https://github.com/zhuang-group/Mesa.

A Survey on Efficient Training of Transformers

“…It demonstrates a promising direction towards more efficient neural scaling laws based on data importance sampling. Rematerialization Herrmann et al [35] Rematerialization ZeRO-Offload [74] Offloading Beaumont et al [7] Offloading + Rematerization ZeRO [72] DP+MP+AMP Megatron-LM [75] DP+TP GPipe [40] DP+PP torchgpipe [48] PP+Rematerization Megatron-LM * [65] DP+TP+PP+AMP Wang et al [84] FP8 Training Cambier et al [11] FP8 Training Mesa [68] 8-bit ACT ACTNN [12], GACT [60] 2-bit ACT [52,42,37] Addition-based PET Bitfit [89], LoRA [38] Reparameterization-based PET…”

“…The saved activations are then dequantized to the original precision in the backward pass to calculate gradients. Recent works [68,60] have been proposed to apply ACT to general frameworks supporting memory-efficient Transformer training.…”

GACT: Activation Compressed Training for Generic Network Architectures

“…"swap" is a simple swapping strategy that swaps all activations to the CPU. For Bert-large, we also show the results on Mesa [7], a memory-saving resource-efficient training framework for transformers, and ZeRO-Offload [37], a highly optimized system for training large-scale language models. Gradient checkpointing uses the default checkpointing policy provided by the transformer library [38], where only the input to each transformer block is saved before the backward pass.…”

“…Although ACT has already demonstrated impressive compression capabilities, previous work on ACT is restricted to specific NN architectures. For example, ActNN [5] is a quantization framework for convolutional NNs only; Mesa [7] proposes a per head/layer quantization method for vision transformers; and AC-GC [6] derives convergence error bound for different types of operators separately.…”

On Block g-Circulant Matrices with Discrete Cosine and Sine Transforms for Transformer-Based Translation Machine

“…As the amount of computational power required for training and inference increases, one of the more advanced neural networks is the Transformer [12], which has a deeper topology. Conversely, the depth of the transformer architecture gives rise to several constraints and challenges, including high computational complexity [13], substantial demands on computational resources [14], and high memory consumption [15] that quadratic to the input sequence length. Therefore, methods are required to achieve excellent performance, mainly when using them as translation machines.…”