Reducing Activation Recomputation in Large Transformer Models

Cut activation memory 5× and recover over 90% of the time lost to recomputation.

Training a transformer stores every intermediate tensor the backward pass will need, and the usual fix when those tensors overflow the GPU is to discard them and recompute them later, at 30 to 40% of the training time. This paper shards the tensors that tensor parallelism copies onto every GPU and recomputes only the attention scores.

For GPT-3 split across 8 GPUs, one layer's stored activations for one sequence fall from 579 MB to 107 MB per GPU, and the 530B-parameter model trains 29.7% faster than it did with full recomputation.

When a large transformer does not fit on its GPUs, the weights are usually not what overflows. The overflow comes from the activations: the intermediate tensors of the forward pass that the backward pass needs to compute gradients. A matrix multiply $Y = XW$ needs $X$ to compute the gradient of $W$, so $X$ has to stay in memory from the moment the forward pass computes it until the backward pass reaches that layer. Every layer of every microbatch in flight holds such tensors at once.

The standard remedy is activation recomputation, also called gradient checkpointing, covered in our gradient checkpointing explainer: keep only each layer's input, drop everything else, and rerun that layer's forward pass when the backward pass needs its internals. The memory drops to almost nothing, and a full extra forward pass is paid for every layer. In NVIDIA's training runs that costs 30 to 40% of the execution time.

This paper, from the team behind Megatron-LM, removes most of that cost with two techniques layered on Megatron's existing tensor parallelism. Sequence parallelism splits, along the sequence, the part of each layer that tensor parallelism leaves copied on every GPU. Selective activation recomputation recomputes only the attention-score tensors, which take up more memory than the rest of a GPT-3 layer and cost little to rebuild. Together they cut activation memory about 5× and recover over 90% of the time full recomputation used to cost. On a 530B-parameter model across 2240 A100 GPUs, MFU (the fraction of the GPUs' peak FLOP rate spent on the FLOPs the model itself needs for its forward and backward passes) rises from 42.1% to 54.2%, a 29% speedup.

The page follows the paper's derivation: count the bytes one layer stores, see which of them tensor parallelism divides and which it copies, split the copied ones along the sequence at the same communication volume, account for what pipeline parallelism adds, and then recompute only the tensors that are large in bytes and small in FLOPs.

Why activations, not weights, fill the GPU

The paper defines an activation as any tensor created in the forward pass that is needed for gradient computation. That excludes the parameters and the optimizer state and includes the dropout masks, since the backward pass must know which entries were zeroed. Weights and optimizer state are fixed once the model and optimizer are chosen. Each activation can either be stored or recomputed, so activations are the part of the memory a training system can trade against compute.

The paper's Figure 1 shows the split for its four test models, 22B to 1T parameters, each on an 80 GB A100. Without the new techniques the activations alone exceed the parameters and optimizer state for every model, by roughly three to four times at 530B and 1T, and the total is above 80 GB in all four cases. With both techniques, every model's total drops below 80 GB.

The other forms of parallelism do not remove the problem. Tensor parallelism (splitting each layer's matrix multiplies across GPUs) does shrink activations, but it communicates inside every layer and its matrix multiplies get smaller as it spreads, so in practice it stays inside one 8-GPU server linked by NVLink. Pipeline parallelism (splitting the layers into consecutive stages on different GPUs) looks as if it should divide activations by the number of stages, and a section below shows why it does not.

Counting one layer's activation bytes

One transformer layer is a layer norm, a self-attention block, a second layer norm and a two-layer MLP, with residual connections around each block. Write $s$ for the sequence length, $b$ for the microbatch size (the number of sequences processed together), $h$ for the hidden width and $a$ for the number of attention heads. The paper assumes 16-bit activations, 2 bytes per element, except dropout masks at 1 byte per element, and ignores small buffers such as a layer norm's per-token mean and variance. A tensor of shape $s \times b \times h$ then costs $2sbh$ bytes.

Operation by operation, each part of the layer keeps what its own backward step reads:

• Each layer norm keeps its input, $2sbh$. Two layer norms: $4sbh$.
• Attention keeps the shared input of the Q, K and V projections ($2sbh$), both Q and K for the $QK^\top$ product ($4sbh$, since the gradient for Q needs K and the one for K needs Q), the softmax output ($2as^2b$, one $s \times s$ matrix per head), the softmax dropout mask ($as^2b$), the dropped-out probabilities and V for the product with the values ($2as^2b + 2sbh$), the input of the output projection ($2sbh$) and the dropout mask after attention ($sbh$). Total: $11sbh + 5as^2b$.
• The MLP keeps the inputs of its two linear layers ($2sbh$ and $8sbh$, the second one $4h$ wide), the GeLU input ($8sbh$) and a dropout mask ($sbh$). Total: $19sbh$.

Adding the three gives $34sbh + 5as^2b$. Factoring out $sbh$ gives the paper's per-layer formula:

The figure lists the fourteen tensors with their sizes for GPT-3 ($a = 96$, $h = 12288$) at one sequence. Click a row to see why the backward pass needs it, and drag the sequence length to watch the three amber score tensors outgrow the eleven teal ones. The two view buttons mark the tensors the next sections are about.

The two terms of Equation (1) grow differently. The $34sbh$ part is linear in the sequence length. The $5as^2b$ part is quadratic in it, because each of the $a$ heads keeps an $s \times s$ score matrix. Written inside the bracket as $5as/h$, it is compared directly with 34: the attention part is larger once $5as/h > 34$, which for GPT-3 happens at $s = 34h/(5a) \approx 870$.

At GPT-3's trained length of 2048, $5as/h$ is exactly 80. With $sbh$ = 2048 × 1 × 12288 = 25,165,824 elements, the structural part is 34 × 25.2M = 856 MB, the attention part is 80 × 25.2M = 2.01 GB, and one layer holds 2.87 GB for a single sequence, 70% of it in the score tensors. GPT-3 has 96 layers, so storing everything for one sequence would take 275 GB.

What tensor parallelism leaves replicated

Megatron's tensor parallelism splits each layer across $t$ GPUs, the tensor-parallel group. In the MLP, the first weight matrix $A$ ($h \times 4h$) is split by columns and the second, $B$ ($4h \times h$), by rows. Each GPU then computes GeLU on its own $4h/t$ columns without talking to the others, because GeLU acts on each element separately, and produces a partial sum of the MLP output. Attention is split the same way across heads. At the end of each block the partial sums are added with an all-reduce, the collective that sums a tensor across all $t$ GPUs and leaves the full sum on every one of them. The paper calls this operator $\bar f$ (all-reduce in the forward pass, nothing in the backward), paired with $f$ at the start of the block (nothing forward, all-reduce backward).

Splitting the matrices also splits the activations inside the blocks: each GPU stores only its heads' score tensors and its $4h/t$ slice of the MLP intermediates. What it does not split is anything at the block boundaries, where every GPU holds the full tensor. That gives the paper's Equation (2):

The 24 is the 34 minus the ten boundary units. Those ten are the two layer-norm inputs ($4sbh$), the inputs to the two blocks, meaning the shared QKV input and the input of the $h \to 4h$ layer ($4sbh$), and the two dropout masks after attention and after the MLP ($2sbh$). Every one of the $t$ GPUs stores an identical copy of them.

For GPT-3 at the paper's $t = 8$, Equation (2) gives $10 + 3 + 10 = 23$ units of $sbh$, or 579 MB per GPU per layer, down from 2.87 GB. The undivided 10 is now 43% of what remains, and it is the same size as the whole attention term. Doubling $t$ to 16 would only bring 23 down to 16.5, and it would push tensor parallelism across the slower links between servers.

Splitting the layer norm along the hidden dimension is possible, and some systems do it (Amazon SageMaker's model parallelism and GSPMD). Then every layer norm must compute each token's mean and variance across devices, an extra collective on an operation that is otherwise cheap.

How sequence parallelism works

Layer norm and dropout treat every token independently. Layer norm normalizes each token's $h$-vector using only that vector; dropout zeroes elements one at a time. So the boundary region can be split along the sequence: rank $i$ keeps tokens $i \cdot s/t$ to $(i+1) \cdot s/t - 1$ and runs layer norm and dropout on those alone. The layer then alternates between two layouts, split by sequence in the layer-norm and dropout regions and split by hidden dimension (as before) inside attention and the MLP, and it needs a conversion at each change.

Going into a block, the matrix multiply needs every token, so the shards are concatenated with an all-gather: each rank sends its shard to all others, and all end with the full tensor. The paper calls this $g$. Coming out, each rank holds a full-length partial sum that must be added across ranks and then cut back into sequence shards. A reduce-scatter does both at once: rank $i$ ends with the sum of shard $i$ only. This is $\bar g$. For the MLP, the paper writes the whole block as Equation (3), with superscripts marking the split dimension:

The same block as code for one rank, with GPT-3's shapes at $t = 8$:

# MLP block with tensor + sequence parallelism, on rank i of t (Eq 3).
# GPT-3 shapes at t = 8: s = 2048, b = 1, h = 12288. Residual add omitted.
def mlp_block(x_i):                  # x_i: [s/t, b, h] = [256, 1, 12288]
    y_i = layernorm(x_i)             # only this sequence shard of Y is stored
    y   = all_gather(y_i)            # g:  [2048, 1, 12288], the full sequence
    z_i = gelu(y @ A_cols[i])        # [2048, 1, 6144], this rank's 4h/t columns
    w_i = z_i @ B_rows[i]            # [2048, 1, 12288], a partial sum
    v_i = reduce_scatter(w_i)        # g-bar: sum over ranks, keep 256 rows
    return dropout(v_i)              # [256, 1, 12288]
# backward: all_gather's gradient is a reduce-scatter and vice versa,
# plus one extra all_gather(y_i) to rebuild Y for the weight gradient of A

In the backward pass each converter runs its mirror image. The gradient of an all-gather is a reduce-scatter, and the gradient of a reduce-scatter is an all-gather, so $g$ is all-gather forward and reduce-scatter backward, and $\bar g$ the reverse.

The new collectives replace the old ones rather than adding to them. Per layer, forward plus backward, tensor parallelism used four all-reduces; tensor plus sequence parallelism uses four all-gathers and four reduce-scatters. A ring all-reduce, the standard bandwidth-optimal implementation, is itself a reduce-scatter followed by an all-gather. Sequence parallelism does the same two halves, with layer norm and dropout run between them on the sequence shards. Step through a ring all-reduce on four ranks in the figure: after the first three steps each rank owns one fully summed sequence chunk, which is exactly the layout $\bar g$ produces, and the last three steps are the all-gather $g$ performs at the start of the next block.

Each rank sends $(t-1)N/t$ bytes in each half, $0.75N$ for $t = 4$, and the pair adds up to the $2(t-1)N/t$ of the all-reduce it replaces. The byte count stays the same, but the timing does not: the paper measures a reduce-scatter plus an all-gather as slightly slower than one all-reduce.

The layer-norm output $Y$ needs extra care. The GeLU matrix multiply needs the full $Y$ again in the backward pass, to compute the gradient of $A$. Storing the gathered $Y$ would put an undivided $2sbh$ back into memory, so each rank stores only its shard $Y_i^s$ and runs one more all-gather in the backward pass, overlapped with the computation of the input gradient so its latency is hidden. That all-gather is additional traffic beyond the swap above; the paper's “no communication overhead” refers to the swap.

With the boundary region sharded, every term carries a $1/t$:

Equation (4) is Equation (1) divided by $t$. For GPT-3 at $t = 8$ it is $114/8 = 14.25$ units, 359 MB per GPU per layer instead of 579 MB. The forward pass also gets a little faster: on one 22B layer it drops from 7.7 ms to 7.2 ms, because layer norm and dropout now run on $1/t$ of the tokens on each GPU.

Why pipeline parallelism does not divide activations

Pipeline parallelism gives each of $p$ GPUs a consecutive group of $L/p$ layers and streams microbatches through them. To keep all stages busy, the 1F1B schedule (one forward, one backward) lets the first stage start $p$ microbatches before the backward pass of the first one returns to it. Until each backward pass finishes, that microbatch's activations stay stored. So the first stage holds $p$ microbatches of $L/p$ layers each, $p \times L/p = L$ layers' worth, whatever $p$ is. The later stages hold fewer, but the first stage sets the memory requirement.

In the figure, the right-hand column counts the microbatches each stage is holding at the time cursor. Play the schedule, then change $p$: stage 0 always peaks at $p$, and the layers-held readout stays at 96.

The total activation memory on the first stage, with tensor and sequence parallelism, is therefore Equation (4) times $L$:

Megatron's interleaved schedule, which gives each GPU $m$ smaller groups of layers to shrink the pipeline's idle time, multiplies this by $1 + (p-1)/(pm)$. The paper also counts the embedding dropout, the final layer norm, the output layer and the fp32 logits; for the 22B model those add less than 0.01%.

Take GPT-3 as the paper runs it: $t = 8$, $p = 8$, $m = 3$, microbatch 1. The interleaving factor is $1 + 7/24 = 1.29$, so the first stage stores about 124 layers' worth. At 579 MB per layer for tensor parallelism alone, that is 71.8 GB of activations. The parameters and optimizer state add about 44 GB per GPU (175B parameters over 64 GPUs is 2.7B each, at 16 bytes per parameter for mixed-precision Adam as counted in the ZeRO explainer), so the total is far past 80 GB. Sequence parallelism alone brings activations to 44.5 GB, and the total to 88.5 GB, still too much.

How selective activation recomputation works

When activations still do not fit, something has to be recomputed. The simplest option is to checkpoint some layers fully and store all activations of the others. It does not scale well: MT-NLG (Megatron-Turing NLG, the paper's 530B model) on 35 pipeline stages has only three layers per GPU, so the choice is between recomputing zero, one, two or three whole layers. The tensors inside a layer also differ widely in what they cost to rebuild.

The $5as/h$ term comes from the operations after the Q, K and V projections widen the tensor to one score matrix per head: the $QK^\top$ product, the softmax, the softmax dropout and the product with V (the region the paper's Figure 3 outlines in red). These tensors are large, but few FLOPs are spent per element. In the paper's FLOP count, the forward pass of a layer spends $24bsh^2$ in the weight matrix multiplies and $4bs^2h$ in the two score products, a ratio of $s : 6h$, 1 to 36 for GPT-3. The $34sbh$ tensors are the opposite case: the GeLU input, for example, comes out of the $h \to 4h$ matrix multiply, one of the most expensive operations in the layer.

Selective activation recomputation therefore stores everything in the 34 (including Q, K and V) and discards only the score tensors. In the backward pass it reruns $QK^\top$, the softmax, the dropout (regenerating the same mask from the saved random-number state) and the product with V from the stored Q, K and V. Memory drops by a fraction of

which is $80/114 = 70\%$ for GPT-3 and $64/98 = 65\%$ for MT-NLG ($a = 128$, $h = 20480$, $5as/h = 64$). The recomputation is one extra forward of the two score products, $2bs^2h$ FLOPs each per layer. The paper puts the resulting ratio of hardware FLOPs (what the GPU executes) to model FLOPs (what the model needs) at approximately $1 + s/6h$, an overhead of 2.7% for GPT-3 and 1.6% for MT-NLG.

The figure shows both bars. Switch between the two models, and drag the sequence length: the memory saving climbs toward 100% as the quadratic term takes over, and the FLOP overhead grows with $s$ as well, since the recomputed products are themselves quadratic in $s$.

In code the change is one checkpoint around the attention core. Megatron-LM exposes it as --recompute-granularity selective, or the shorthand --recompute-activations.

# Selective activation recomputation: checkpoint only the attention core.
def attn_core(q, k, v):                  # q, k, v: [b, a, s, h/a], stored
    p = softmax(q @ k.transpose(-1, -2)) # [b, a, s, s]: 2as^2b bytes
    p = dropout(p)                       # mask as^2b, output 2as^2b
    return p @ v                         # [b, a, s, h/a]

def layer(x):                            # x: [s, b, h]
    y = layernorm(x)
    q, k, v = qkv(y)                     # kept: part of the 34sbh
    o = checkpoint(attn_core, q, k, v)   # scores dropped, rebuilt in backward
    x = x + dropout(proj(o))
    return x + dropout(mlp(layernorm(x)))

With sequence parallelism and selective recomputation together, the first stage stores

(multiplied by $1 + (p-1)/(pm)$ under interleaving). The attention term is gone, so activation memory is now linear in $s$ and does not depend on the number of heads. For GPT-3 that is $34/8 = 4.25$ units, 107 MB per layer and 13.3 GB for the 124 layers' worth the first stage holds. With 44 GB of parameters and optimizer state the total is about 57 GB, inside 80 GB with no layer fully recomputed. Selective recomputation without sequence parallelism would leave $10 + 24/8 = 13$ units, 40.6 GB of activations, and would not fit; the paper reports that none of its four models fit without both techniques together.

What the paper measured

All experiments use $t = 8$, $s = 2048$, a vocabulary of 51,200 and no data parallelism, on 80 GB A100s. The four models are 22B ($L = 48$, $p = 1$, 8 GPUs), 175B GPT-3 ($L = 96$, $p = 8$, 64 GPUs), 530B MT-NLG ($L = 105$, $p = 35$, 280 GPUs) and 1T ($L = 128$, $p = 64$, 512 GPUs), with the interleaved schedule ($m = 3$) for 175B and 530B.

Memory first, as a percentage of the tensor-parallel baseline of Equation (2). The paper reports that each technique alone cuts the requirement nearly in half (by the Table 2 formulas, to 62% and 57% for GPT-3), and that together they bring it under 20%, the 5× reduction of the title. Full recomputation reaches about 10%, so the combination uses roughly twice the memory of full recomputation. For GPT-3 the formulas give 23 units for the baseline, 4.25 for both techniques (18.5%, a 5.4× cut) and 2 for full recomputation (8.7%).

On one layer of the 22B model (Table 4), the forward plus backward pass takes 19.6 ms with no recomputation. Full recomputation raises it to 27.2 ms, 39% more; the expected cost of one extra forward is about 33%, and the rest comes from losing an optimization that overlaps the backward all-reduce with the weight gradient computation. Selective recomputation costs 20.9 ms (7%), and selective recomputation with sequence parallelism 20.3 ms (4%). On the larger models the per-layer overhead falls to 2% for 530B and 1T, against 36% for full recomputation.

End to end (Table 5), the full training iteration runs 29.0% to 32.1% faster than with full recomputation, taking about 76 to 77% of the time. For 530B the iteration drops from 49.05 s to 37.83 s. MFU rises with model size, to 56.3% for the 1T model. HFU, the same ratio computed with the recomputed FLOPs included, is 57.0%. Both are measured against the A100's 312 TFLOP/s dense 16-bit peak.

The abstract's 54.2% is a separate run: the 530B model with 8-way data parallelism on 2240 GPUs (280 × 8), where the gradient all-reduce between replicas raises the iteration time from 37.83 s to 39.15 s and MFU drops from Table 5's 56.0% to 54.2%. The paper did not overlap that all-reduce with the backward pass. The 42.1% it compares against is the same setup with full recomputation, and 54.2/42.1 is the 29% speedup.

Scope and later work

The evaluation runs without data parallelism apart from that one 2240-GPU data point, so batch sizes (4 to 512 sequences) are far below real training. Data-parallel memory techniques such as ZeRO and CPU offloading are described as complementary and are not compared. All memory numbers include an optimization from Appendix B that frees each pipeline stage's output tensor after its forward pass, since the next stage holds a copy. Appendix C tests storing all activations for as many microbatches as memory allows and recomputing only the rest, and finds that after the two main techniques it adds little.

The attention score tensors were also the target of FlashAttention, released later the same month. It computes attention in tiles, keeps only the output and the softmax normalization statistics, and recomputes the scores inside one fused kernel in the backward pass, which removes the $5as^2b$ term without a separate checkpoint. Sequence parallelism addresses the replicated $10sbh$, which FlashAttention does not touch, and it remains part of Megatron-LM.