Tempo: Accelerating Transformer-Based Model Training through Memory Footprint Reduction

TL;DR

Tempo is a memory-efficient method for training Transformer models that enables larger batch sizes and significantly improves training throughput without sacrificing accuracy.

Contribution

Tempo introduces drop-in memory reduction techniques for key Transformer layers, enabling faster training and larger batch sizes on existing hardware.

Findings

Up to 2x larger batch sizes with Tempo.

16% higher training throughput on BERT Large.

19-26% speedup on GPT2 and RoBERTa models.

Abstract

Training deep learning models can be computationally expensive. Prior works have shown that increasing the batch size can potentially lead to better overall throughput. However, the batch size is frequently limited by the accelerator memory capacity due to the activations/feature maps stored for the training backward pass, as larger batch sizes require larger feature maps to be stored. Transformer-based models, which have recently seen a surge in popularity due to their good performance and applicability to a variety of tasks, have a similar problem. To remedy this issue, we propose Tempo, a new approach to efficiently use accelerator (e.g., GPU) memory resources for training Transformer-based models. Our approach provides drop-in replacements for the GELU, LayerNorm, and Attention layers, reducing the memory usage and ultimately leading to more efficient training. We implement Tempo and evaluate the throughput, memory usage, and accuracy/loss on the BERT Large pre-training task. We demonstrate that Tempo enables up to 2x higher batch sizes and 16% higher training throughput over the state-of-the-art baseline. We also evaluate Tempo on GPT2 and RoBERTa models, showing 19% and 26% speedup over the baseline.