Reducing Activation Recomputation in Large Transformer Models

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 5×, while reducing execution time overhead from activation recomputation by over 90%. For example, when training a 530B parameter GPT-3 style model [20] 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 1 and NeMo-Megatron 2 .

Fold3D: Rethinking and Parallelizing Computational and Communicational Tasks in the Training of Large DNN Models

“…Baselines. We took Megatron v3.0 (Megatron-SP) [28], Megatron v2.5 (Megatron-PTD) [8], DeepSpeed 3D (DSpeed3D) [19], and DeepSpeed ZeRO3 (DSpeedZ3) [29] as our baselines. Megatron-SP is the latest 3D parallel training system that was reported to achieve almost linear scaling efficiency.…”

“…We ran these two DeepSpeed systems in DeepSpeed v0.5.5 environment. Sequence parallelism [28] was integrated into Megatron, DSpeed3D and FOLD3D to reduce the activation size and support larger models.…”

“…We compared FOLD3D against Megatron-SP [28] (v3.0.2, the latest release), Megatron-PTD [8] (v2.5.0), DeepSpeed Zero3 [29] (DSpeedZ3), and DeepSpeed 3D [19] (DSpeed3D), covering three notable and open source 3D training systems and one state-of-the-art data parallel training system (DSpeedZ3). Both FOLD3D and Megatron-SP are enabled with sequence parallelism (Section V), one of the latest memory squeezing techniques [28] complementary to 3D parallelism. Our evaluation was done both over a high-profile cluster (256 A100 GPUs) and a middle-profile cluster (64 V100 GPUs).…”

SuperScaler: Supporting Flexible DNN Parallelization via a Unified Abstraction

“…Inspired by Einops [41], we apply annotations on tensor dimensions to indicate transformation strategies. The Categories Mechanisms Support SPMD Parallelism Data Parallelism [1] Sequence Parallelism [24] Transformer Parallelism [45] DAP [11] ZeRO [38] Sequence Parallelism [26] * Flexible Tensor Parallel [20,53,56] MPMD Parallelism 1F1B [45,50] GPipe [19] Chimera [27] PipeDream (Async) [33] × Terapipe [28] ×…”

“…In addition, various memory optimizations [10,18,23] have been adopted to exploit large-scale model training under GPU memory constraints. Recently, systems, e.g., Megatron-LM [24,34,45], DeepSpeed [38], Piper [46], Unity [51] and Alpa [61], combine multiple parallelisms and memory optimizations within one system to accelerate distributed DNN training. However, these solutions fall short in relying on empirical parallelism configurations and having limited execution scheduling choices.…”