Reducing Activation Recomputation in Large Transformer Models

TL;DR

This paper introduces simple techniques to significantly reduce activation memory and recomputation overhead in large transformer training, enabling faster and more efficient scaling to models with up to one trillion parameters.

Contribution

The authors propose sequence parallelism and selective activation recomputation, reducing activation memory by 5x and recomputation overhead by over 90%, improving large transformer training efficiency.

Findings

Activation memory reduced by 5x.

Recomputation overhead decreased by over 90%.

Training a 530B parameter model is 29% faster.

Abstract

Training large transformer models is one of the most important computational challenges of modern AI. In this paper, we show how to significantly accelerate training of large transformer models by reducing activation recomputation. Activation recomputation is commonly used to work around memory capacity constraints. Rather than storing activations for backpropagation, they are traditionally recomputed, which saves memory but adds redundant compute. In this work, we show most of this redundant compute is unnecessary because we can reduce memory consumption sufficiently without it. We present two novel yet very simple techniques: sequence parallelism and selective activation recomputation. In conjunction with tensor parallelism, these techniques almost eliminate the need to recompute activations. We evaluate our approach on language models up to one trillion parameters in scale and show that our method reduces activation memory by 5x, while reducing execution time overhead from activation recomputation by over 90%. For example, when training a 530B parameter GPT-3 style model on 2240 NVIDIA A100 GPUs, we achieve a Model Flops Utilization of 54.2%, which is 29% faster than the 42.1% we achieve using recomputation. Our implementation will be available in both Megatron-LM and NeMo-Megatron.