激活重新计算权衡

Training a neural network with backpropagation has two phases. In the forward pass, the input flows through each layer and every layer produces an output, called an activation. In the backward pass, the training algorithm works from the loss back toward the input, computing a gradient for every weight. Most of those gradient calculations need the activations from the forward pass, so by default a framework such as PyTorch or JAX keeps every activation in memory until the backward pass uses it. Activation recomputation, also known as activation checkpointing, gradient checkpointing or rematerialization, changes that rule: it keeps only some activations, discards the rest, and recomputes the discarded ones from the nearest kept value when the backward pass needs them.

The reason this matters is that activation memory, unlike weight memory, grows with the size of each training batch and with sequence length. For large transformers it often dominates. Korthikanti and colleagues at NVIDIA give a formula for a single transformer layer with no model parallelism: activation memory in bytes is roughly s times b times h times (34 + 5as/h), where s is sequence length, b is microbatch size, h is hidden size and a is the number of attention heads. The 5as/h part comes from attention, which builds a score for every pair of tokens, so it grows with the square of the sequence length. For a GPT-3-sized configuration that term equals 80, larger than the 34 from everything else. Multiplied across nearly a hundred layers, the stored activations can exceed the memory of any single GPU even before the weights are counted.

The mechanism is simple to state. Pick a set of checkpoints, usually the inputs to certain layers or blocks. During the forward pass, save only those checkpoints and throw away everything computed between them. During the backward pass, when gradients reach a block, rerun that block's forward computation from its saved input to rebuild its internal activations, use them to compute gradients, and free them again before moving to the previous block. Only one block's worth of full activations needs to be in memory at a time, plus the checkpoints. The price is that most forward computation is done twice.

The idea was brought into mainstream deep learning practice by Tianqi Chen, Bing Xu, Chiyuan Zhang and Carlos Guestrin in the April 2016 paper Training Deep Nets with Sublinear Memory Cost. They showed that for a network of n layers, placing a checkpoint every square root of n layers cuts activation memory from proportional to n to proportional to the square root of n, at the cost of roughly one extra forward pass per mini-batch. They also showed a more extreme version that reduces memory to proportional to log n at the cost of order n log n extra forward computation. Their headline example reduced the memory for a 1,000-layer residual network from 48 GB to 7 GB with about 30 percent more running time. In 2019, the Checkmate paper by Paras Jain and colleagues at Berkeley, published at MLSys 2020, framed checkpoint selection as a mixed integer linear program that finds an optimal recomputation schedule for a given memory budget in under an hour with off-the-shelf solvers, and reported training with up to 5.1 times larger inputs.

For large language models, the key refinement is selective recomputation. Korthikanti and colleagues observed in May 2022 that not all activations are equally expensive to recompute. The attention operations, meaning the query-key product, the softmax, softmax dropout and the weighted sum over values, create large activations but require relatively few floating point operations per element. The large matrix multiplications in the feed-forward layers are the opposite. By recomputing only the attention pieces and storing everything else, combined with sequence parallelism, which splits certain layer operations across GPUs along the sequence dimension, they reduced activation memory by about five times and cut the time overhead of recomputation by more than 90 percent. For a single 22-billion-parameter layer, they measured full recomputation adding 39 percent execution time and the combined approach only 4 percent. End to end, training throughput rose by about 30 percent across models from 22 billion to 1 trillion parameters.

Recomputation competes with several alternatives, and real systems mix them. Doing nothing is fastest when memory allows. Shrinking the microbatch reduces activation memory linearly but hurts GPU efficiency. Offloading moves activations to CPU memory and back, trading compute for data transfer; PyTorch's selective checkpointing API and JAX's policies both include offload options. Model parallelism splits activations across GPUs at the cost of communication. Fused kernels such as FlashAttention, posted in May 2022 by Tri Dao and colleagues, never write the full attention matrix to GPU memory: they store the softmax normalization statistics and recompute attention blocks on-chip during the backward pass. The paper reports this makes memory linear in sequence length and speeds up the backward pass despite the extra arithmetic, because reading and writing GPU high-bandwidth memory was the real bottleneck.

The limitations are mostly about cost and correctness. Full recomputation adds roughly 30 to 40 percent of execution time in the NVIDIA measurements, which translates directly into training cost. Placing checkpoints poorly saves little memory: the JAX documentation notes that checkpointing the entire differentiated function, or only its last stage, saves nothing and wastes compute. Recomputation must also reproduce the forward pass exactly. Layers that use randomness, such as dropout, need the random number generator state restored, which PyTorch does by default through preserve_rng_state at some performance cost. The PyTorch documentation warns that if a checkpointed function behaves differently on the second run, for example because it reads a global variable, the result can be an error or silently incorrect gradients.

In practice, using recomputation requires little code but careful measurement. In PyTorch, wrapping each transformer block with torch.utils.checkpoint.checkpoint and passing use_reentrant=False, which the documentation recommends, is the usual starting point; selective policies can then save specific operations such as matrix multiplies. In JAX, jax.checkpoint with a policy does the same, and the documentation notes that under jax.jit the XLA compiler often makes these choices automatically. The right setting is found by profiling peak memory and step time together: the goal is the largest batch or sequence length that fits, at the lowest recomputation overhead, since every percent of overhead is paid on every training step.