Estimating Activation Memory in Training

When people first estimate how much memory a model needs, the model's weights usually receive most of the attention. The weights are the learned parameters, and it is natural to assume they are the largest consumer of memory. For large models, that assumption is often wrong. During training, a second group of values, called activations, can take up a surprising amount of memory, and in some cases far more than the weights themselves.

This article looks at why that happens and at one of the main techniques used to reduce it. Keep one question in mind as you read: if activations are stored so that gradients can be computed later, what is one way to shrink the memory they occupy without changing the model itself?

The weights have a fixed footprint set by the network's architecture. Once the model is defined, the number of parameters does not change from one batch to the next. Activations behave differently. They are the intermediate values produced as data moves forward through the network, and when they are kept so that the gradient computation can use them later, they accumulate. The stored activations depend on the work the network does on each batch as it moves forward, which is why their combined size can grow to dwarf the memory taken by the weights.

The diagram shows the two groups side by side. The weights form one block with a fixed size. The activations form a growing stack, one entry for each layer whose output is kept. Adding more layers, or processing larger batches, adds more entries to the activation stack while the weight block stays the same.

Published measurements show the effect at real scale. A chart compares the memory needed for activations with the memory needed for the weights across several Megatron models of different sizes. The numbers come from work by Korthikanti and colleagues (2022). The main lesson is visible across the scales: as models grow, the activation side of the comparison can become the larger share.

If stored activations are the problem, one direct response is to stop storing all of them. Instead of keeping activations around for reuse, the training process recomputes them when they are needed. This technique is called gradient checkpointing, and it is also known as activation recomputation. The two names refer to the same idea, and you will see both in the literature.

Follow one activation through its life under this technique. It is produced during the forward pass. It is not kept in memory for the whole run. When the gradient step needs it, it is produced again. The diagram below shows the order of these events. It is a simplified view of the sequence, not a description of how any particular framework schedules its work.

The core move is a swap of storage for computation. Memory is saved because activations are not kept between the forward pass and the point where the gradient calculation needs them. Work is added because those activations must be produced again at that later point. The technique reduces memory requirements, but it increases the time needed for training because of the recomputation. Neither side of this exchange is free.

Which strategy to use depends on which constraint is tighter. When memory is the constraint that prevents a model from being trained at all, the extra time may be a reasonable price. When time is the tighter constraint and memory already fits, storing activations as usual may be the better choice. The source names both effects and does not claim that one always beats the other.

Activation memory is one of the main costs of training, and it can exceed the memory used by the weights in large models. Gradient checkpointing, also called activation recomputation, reduces that cost by recomputing activations when they are needed rather than storing them. The reduction in memory is paid for with additional training time, so the choice is a deliberate trade.