Reducing Activation Recomputation in Large Transformer Models

By Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence C. McAfee, Michael Andersch, Mohammad Shoeybi, Bryan Catanzaro. 2022

492 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Efficiency & Compression

Citations in the corpus

LLaMA: Open and Efficient Foundation Language Models (2023, methodology): “This is achieved by manually implementing the backward function for the transformer layers, instead of relying on the PyTorch autograd. reduce the memory usage of the model by using model and sequence parallelism, as described by Korthikanti et al. (2022).”
PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, background): “For example, sequence parallelism [12] reduces activation memory in conjunction with tensor parallelism; Pipetransformer [5] designed a dynamic 2D parallelism that allows changing the dimensions of pipeline and data parallelism on the fly, depending on learning signals.”
Efficiently Scaling Transformer Inference (2022, neutral): “Prior works propose several approaches for efﬁcient partitioning to train large models efﬁ-ciently, for e.g., NeMo Megatron (Korthikanti et al., 2022), GSPMD (Xu et al., 2021) and Alpa (Zheng et al., 2022).”
A Survey of Large Language Models for Healthcare: from Data, Technology, and Applications to Accountability and Ethics (2023, unknown): “including Megatron-LM [35], OPT [36], Chinchilla [37], Galactica [38], LLaMA [39], PaLM [40], and PaLM-2 [16].”
DeepSeek-VL: Towards Real-World Vision-Language Understanding (2024, methodology): “With such minor modification, we can now perform canonical 3D parallelism techniques as in Megatron (Korthikanti et al., 2023; Narayanan et al., 2021; Shoeybi et al., 2019) and overlap computation and communication as in DeepSeek-LLM (DeepSeek-AI, 2024).”
Ring Attention with Blockwise Transformers for Near-Infinite Context (2023, methodology): “We suggest considering porting our method to CUDA, OpenAI Triton, or Jax Pallas in the future, for both maximum sequence length and maximum compute performance.”
HunyuanVideo: A Systematic Framework For Large Video Generative Models (2024, unknown): “HunyuanVideo training adopts 5D parallel strategies, including tensor parallelism (TP) [74], sequence parallelism (SP) [45], context parallelism (CP) [63], and data parallelism combined with Zero optimization (DP + ZeroCache [62]).”
MegaScale: Scaling Large Language Model Training to More Than 10, 000 GPUs (2024, methodology): “We also evaluate the effect of LAMB optimizer as depicted in Figure 10b, which shows that LAMB optimizer with four times of batch size achieves the same loss as ADAM optimizer after around 250B tokens.”
A Survey on Mixture of Experts in Large Language Models (2024, unknown): “…[34], [105], [108], [163], [164] with other existing parallel strategies (tensor [165], [166], [167], pipeline [168], [169], [170], sequence parallelism [171], [172], [173]) has been investigated to enhance the scalability and efﬁciency of MoE models in large-scale distributed environments.”
DeepSeek-Coder: When the Large Language Model Meets Programming - The Rise of Code Intelligence (2024, methodology): “These include tensor parallelism (Korthikanti et al., 2023), alongside ZeRO data parallelism (Rajbhandari et al., 2020) and PipeDream pipeline parallelism (Narayanan et al., 2019 utilize clusters outfitted with NVIDIA A100 and H800 GPUs.”
ASPEN: High-Throughput LoRA Fine-Tuning of Large Language Models with a Single GPU (2023, unknown): “Vijay et al. [25] establish the functional relationship between memory size and input training data size.”
AI and Memory Wall (2024, unknown): “The rest of the activations could then be recomputed when needed [15, 19].”
DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models (2024, unknown): “× We first compare DeepSeekMoE with a larger GShard model with 1.5 times the expert size, which results in 1.”
Qwen-Image Technical Report (2025, unknown): “Therefore, we leverage Megatron-LM (Shoeybi et al., 2019; Korthikanti et al., 2023) for training and apply the following optimizations: Hybrid Parallelism Strategy We adopted a hybrid parallelism strategy, combining data parallelism and tensor parallelism, to efficiently scale training across large…”
LongVILA: Scaling Long-Context Visual Language Models for Long Videos (2024, methodology): “Additionally, we compare our sys-tem with the expert-designed and highly optimized Megatron-LM(Shoeybi et al., 2019; Korthikanti et al., 2023) system, focusing on their implementation of sequence parallelism, termed “context parallelism” (CP).”
BloombergGPT: A Large Language Model for Finance (2023, unknown): “After GPT-3, models grew in size to 280 billion (Gopher, Rae et al., 2021), 540 billion (PaLM, Chowdhery et al., 2022), and 1 trillion parameters (Megatron, Korthikanti et al., 2022).”
FLUX.1 Kontext: Flow Matching for In-Context Image Generation and Editing in Latent Space (2025, unknown): “We use selective activation checkpointing [26] to reduce maximum VRAM usage.”