PaperlayerView exact revision on arXiv

Source-linked AI summary

Reducing Activation Recomputation in Large Transformer Models

Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, Bryan Catanzaro

arXiv:2205.05198v1cs.LGcs.CL

TL;DR

Large transformer training is constrained by activation memory, while full activation recomputation saves memory at a 30 −40% execution-time overhead. The paper combines sequence parallelism with tensor parallelism and selective activation recomputation, achieving a 5× memory reduction and recovering over 90% of that recomputation overhead.

  • Problem

    Activation memory can exceed GPU capacity for very large transformers, so training commonly uses full activation recomputation, which adds substantial execution overhead.

  • Method

    The paper combines sequence parallelism with tensor parallelism and selectively recomputes large, inexpensive-to-recompute attention activations.

  • Results

    5× memory reduction and over 90% recovery of the compute overhead introduced by full activation recomputation are achieved.

  • Takeaways & Limitations

    The techniques substantially reduce activation-memory pressure and the need for activation recomputation while retaining compute efficiency.

  • Takeaways & Limitations

    Future work targets memory fragmentation for large microbatches, non-uniform pipeline allocation, and pressure on the first pipeline stage.

Abstract

from arXiv · show

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.

1 Introduction

Training very large transformers requires distributing model state across GPUs, but activation memory can exceed device capacity and full recomputation adds substantial execution overhead. The paper introduces simple techniques that reduce activation-memory pressure and recomputation cost.

  • Model parallelism distributes parameters, activations, and optimizer state across devices for trillion-parameter training.Tensor-level parallelism faces communication and matrix-multiplication efficiency limits as the number of devices grows.
  • Baseline memory requirements for models from 22B to 1T parameters exceed the 80GB capacity of an NVIDIA A100 GPU.The standard response is to avoid storing most activations and recompute them during backpropagation.
  • 30 −40% execution time overhead is observed when full activation recomputation is used.Prior transformer training commonly checkpoints activations at layer boundaries and recomputes the remaining activations during the backward pass.
  • The paper presents transformer-specific techniques that reduce activation-memory pressure with no, or very low, impact on compute efficiency.The techniques are positioned as complementary to data partitioning and CPU offloading, which have higher implementation cost and larger compute-efficiency impacts.
  • Sequence parallelism and selective activation recomputation are introduced to reduce redundant activation storage and recomputation.The paper first develops an activation-memory formula, then studies model-parallelism effects and selectively recomputes expensive-to-store, inexpensive-to-recompute activations.

2 Related Work

Related work distributes model state through tensor, pipeline, or data parallelism, with different trade-offs in scalability, communication, memory replication, and compute efficiency.

  • Tensor parallelism distributes each layer’s parameters across devices, while pipeline parallelism splits the network along its layer dimension.Combining both forms has enabled training models up to 1T parameters.
  • Model parallelism enables large-model training because parameters and optimizer states do not fit on a single GPU.Even host-device swapping can leave compute requirements unrealistically long, motivating parallelism.
  • Data-parallel approaches shard optimizer states, gradients, and parameters across data-parallel ranks, with some extensions offloading data to CPU memory.These methods provide an alternative to model parallelism for large-scale training.
  • Sequence parallelism partitions activations along sequence dimensions throughout the network to alleviate regions where tensor parallelism leaves activations unsplit.The cited prior approach replicates parameters and optimizer state across devices, limiting its suitability for large models.

3 Transformer Architecture

The paper analyzes a single-stack transformer encoder or decoder with layered self-attention and MLP blocks, using sequence length, microbatch size, hidden dimension, and vocabulary size to describe tensor shapes.

  • The model is a single stack of L transformer layers, each receiving and producing a tensor of size s × b × h.Here s is sequence length, b is microbatch size, and h is hidden dimension.
  • Each transformer layer combines self-attention with a two-layer MLP that expands the hidden size to 4h and reduces it back to h.The self-attention block contains a attention heads.
  • Input tokens are embedded using a v×h word-embedding table and s×h learned positional embeddings.The embedding output is the s × b × h input to the transformer stack, and word-embedding and output-layer weights are shared.
  • The final transformer output is projected into the vocabulary dimension to calculate cross-entropy loss.

4 Activation Memory

This section quantifies transformer activation memory and shows how tensor plus sequence parallelism reduces storage across parallel ranks. Pipeline scheduling and selective recomputation further motivate approximating and reducing the activation-memory burden.

  • Activation-memory accounting: Activations are forward-pass tensors needed for gradient computation, excluding model parameters and optimizer state but including dropout masks.The analysis uses 16-bit storage for most elements and one byte per dropout-mask element.
  • Activation-memory accounting: The attention block requires 11sbh + 5as^2b bytes, while the MLP block requires 19sbh bytes of activation storage.These totals combine the principal stored inputs, attention intermediates, dropout masks, and nonlinear-operation inputs.
  • Tensor parallelism: Tensor parallelism divides activations inside attention and MLP blocks, but replicated layer-normalization and dropout operations leave the 10sbh term undivided.The computationally intensive parts are parallelized, whereas replicated operations contribute substantial activation memory without requiring much compute.
  • Sequence parallelism: Sequence parallelism partitions non-tensor-parallel regions along the sequence dimension and introduces communication converters between sequence- and tensor-parallel regions.The combined approach uses g and ¯g operations to coordinate all-gather and reduce-scatter communication around these regions.
  • Sequence parallelism: Using tensor and sequence parallelism distributes activations across the tensor-parallel group and reduces required per-layer activation memory by tensor-parallel size t.The resulting per-layer expression is the tensor-parallel expression divided by t.
  • Pipeline parallelism: Under the 1F1B pipeline schedule, the first pipeline stage stores activations for p microbatches, totaling L layers worth of activations regardless of pipeline size p.For a 22B-parameter model, extra terms beyond the main approximation account for less than 0.01% of total activation memory.

5 Selective Activation Recomputation

Selective activation recomputation stores and recomputes only memory-intensive, inexpensive-to-recompute parts of transformer layers, reducing activation memory with modest computational cost.

  • Motivation: Full activation recomputation can reduce memory substantially but introduces as much as 30 −40% computational time overhead.Checkpointing only enough activations for the model-parallel configuration balances memory capacity and recomputation cost.
  • Interaction with sequence parallelism: Sequence parallelism enables more configurations by reducing activation memory, allowing selective recomputation instead of recomputing full transformer layers.An alternative that stores only portions of activations per tensor-parallel rank requires an extra all-gather per layer and is not considered.
  • Selective activation recomputation: Selective activation recomputation checkpoints only memory-intensive transformer operations whose recomputation requires relatively few FLOPs.The targeted operations include QKT multiplication, softmax, dropout, and attention over V.
  • Selective activation recomputation: 70% memory savings for GPT-3 and 65% for MT-NLG come with only 2.7% and 1.6% FLOPs overhead, respectively.For GPT-3, 5as/h = 80; for MT-NLG, 5as/h = 64, exceeding the factor 34 for the rest of the layer.
  • Memory scaling: With selective activation recomputation, activation memory scales linearly with sequence length and is independent of the number of attention heads.For an interleaved pipeline schedule, the resulting memory expression is multiplied by (1 + p−1/m).

6 Evaluations

Evaluations on models up to one trillion parameters show that sequence parallelism and selective recomputation substantially reduce memory and recomputation overhead while improving throughput and utilization.

  • 6.1 Memory Usage: 5x combined memory reduction brings activation requirements below 20% of the tensor-parallel baseline, compared with 10% for full recomputation.Each technique individually nearly halves memory, and the combined approach is about 2x the full-recomputation memory.
  • 6.2 Execution Time per Layer: 4% combined overhead is achieved with sequence parallelism and selective recomputation, versus 39% for full-layer recomputation.Selective recomputation alone has 7% combined forward-and-backward overhead, compared with 39% for full recomputation.
  • 6.2 Execution Time per Layer: 2% overhead for the 530B and 1T cases compares with 36% overhead for full recomputation as model size increases.Figure 8 reports the per-layer forward, backward, and recomputation breakdown across test cases.
  • 6.3 End-to-End Iteration Time: 29.0% to 32.1% throughput improvement is obtained across all tested configurations over full recomputation without sequence parallelism.The reported end-to-end iteration results use the configurations listed in Table 3 and translate into shorter training times.
  • 6.3 End-to-End Iteration Time: 56.3% MFU and 57.0% HFU are reached for the one-trillion-parameter model.MFU and HFU are based on model and hardware FLOPs per second divided by accelerator theoretical peak FLOPs per second.
  • 6.3 End-to-End Iteration Time: 54.2% MFU results when scaling the 530B model to 8-way data parallelism across 2240 GPUs, down from 56.0% without that scaling.Iteration time increases from 37.83 seconds to 39.15 seconds under this configuration.

7 Conclusions and Future Work

The paper concludes that sequence parallelism combined with selective activation recomputation reduces activation memory and recovers most of full recomputation’s compute overhead, while identifying remaining pipeline-memory challenges.

  • Conclusions: 5× lower activation memory and over 90% recovery of full activation recomputation’s compute overhead are achieved with the proposed techniques.The techniques reduce memory pressure from storing activations and thereby reduce the need for recomputation.
  • Future Work: Future work targets memory fragmentation, non-uniform allocation from pipeline parallelism, and memory pressure on the pipeline’s first stage.These issues are identified as opportunities for further activation-memory reduction.

A FLOPs Calculation

The paper calculates model and hardware FLOPs by summing the dominant transformer matrix multiplications, then compares them to show that selective recomputation adds little overhead.

  • Attention and feed-forward computation: 6Bsh^2, 2Bs^2h, 2Bs^2h, and 2Bsh^2 operations account for the attention block’s main forward-pass FLOPs.These terms represent key, query, and value transformations, attention matrix computation, attention over values, and post-attention projection.
  • Attention and feed-forward computation: 16Bsh^2 FLOPs arise from the feed-forward network, producing 24Bsh^2 + 4Bs^2h FLOPs per transformer-layer forward pass.The feed-forward network expands the hidden size to 4h and reduces it back to h.
  • Language-model head: 2Bshv FLOPs are required for the language-model logits layer, which transforms features of dimension h to vocabulary dimension v.
  • Forward and backward totals: The backward pass requires double the forward-pass FLOPs because gradients are computed for both input and weight tensors.The paper then sums these contributions to define model FLOPs for one forward and backward pass.
  • Recomputation overhead: Selective activation recomputation adds 2Bs^2h for attention-matrix computation and 2Bs^2h for attention over values to the per-iteration hardware FLOPs.These additional operations are included in the hardware-FLOPs total.
  • Recomputation overhead: The hardware-to-model FLOPs ratio is close to one under the assumptions 3h ≫ s and 12hL ≫ v, indicating little extra computation.

B Pipeline Parallelism Memory Optimization

Pipeline-parallel activation memory is imbalanced across ranks, and deallocating redundant output tensors reduces that memory, especially on the first pipeline stage.

  • Memory imbalance: The 530B-model experiment compares activation memory across pipeline ranks before and after deallocating each rank’s output tensor.The unoptimized case is shown in blue and the memory-optimized case in yellow.
  • Memory imbalance: Activation memory decreases linearly along pipeline ranks, while rank 0 has an additional sbhp spike from the embedding layer.
  • Output-tensor deallocation: The optimization saves sbhr memory per pipeline rank, peaking at r = p on the first pipeline stage.Here r is the number of microbatches in flight on each rank, and p is the pipeline-parallel size.
  • Output-tensor deallocation: 2.73 GB of theoretical savings on the first pipeline stage closely matches the rank-0 difference between the plotted memory lines.The estimate uses the 530B model’s hyperparameters and 2 bytes per data element.

C Microbatch Level Activation Recomputation

Microbatch-level recomputation uses available memory to store complete activations for selected microbatches, reducing checkpointing as backpropagation frees space and improving utilization modestly.

  • Memory pressure: The first pipeline stage stores the most activations, equivalent to storing activations for all transformer layers in the model.
  • Execution pattern: Figure 10 compares forward, recomputation, and backpropagation patterns for the baseline and microbatch-level strategies.Yellow boxes denote forward passes with checkpointed activations, red boxes recomputation, blue boxes backpropagation, and white boxes forward passes with all activations saved.
  • Moving activation window: Microbatch-level recomputation stores all activations for outstanding microbatches until device memory is used, then checkpoints the remaining microbatches.Backpropagation later frees enough memory to store another full layer of activations.
  • Stage-dependent recomputation: Later pipeline stages have fewer outstanding backpropagation steps, and many therefore require no activation recomputation.Outstanding steps at stage S are calculated as max(0, p − S).
  • Utilization result: 52.3% (+0.7%) MFU is reached for the 175B model, compared with the baseline using sequence parallelism and selective activation recomputation.
  • Utilization result: 56.4% (+0.4%) MFU is reached for the 530B model, while selective recomputation overhead is approximately 2%.The reported gain is small because selective recomputation already has low overhead.
Loading 2205.05198v1…