Trade-offs in Training Large Models
Activation memory can be much larger than the memory used by model weights, so the weights are not always the main memory cost.
The Weights Are Not Always the Main Cost
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.
What do you think happens?
During training of a large model, which takes up more memory: the learned weights, or the activations produced during the forward pass?
Reveal answer
For large models, the activations can be much larger than the memory used by the weights, so the weights are not always the main memory cost.
Weights have a fixed footprint set by the architecture. Activations accumulate with the work the network does on each batch as it moves forward, so their total can grow to dwarf the weights.
This article looks at why activation memory grows so large, and at one of the main techniques used to reduce it: gradient checkpointing. That technique does not change the model itself. It changes what gets stored, and it pays for the saved memory with extra computation.
Where Activation Memory Comes From
Activations are the values a network produces as data moves forward through its layers. If the activations produced during the forward pass are stored so that the gradient computation can use them later, their combined size can grow very large. The reason lies in when activations are kept, not in the model's design alone.
The weights have a fixed footprint, because the architecture determines how many there are. The stored activations, by contrast, depend on the work the network does on each batch as it moves forward, so they keep accumulating. As models grow, the activation side of the comparison can become the larger share. Published measurements comparing these two quantities across several Megatron models at different scales were reported by Korthikanti and colleagues (2022). The exact values should be read from that original paper rather than reconstructed from memory.
Gradient Checkpointing as a Storage-for-Computation Swap
Gradient checkpointing, also called activation recomputation, reduces memory by not storing activations for reuse and recomputing them when they are needed. The two names refer to the same idea, and you will see both in the literature.
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 therefore reduces memory requirements, but it increases the time needed for training because of the recomputation. Neither side of the exchange is free.
Following One Activation Through Its Life
Trace one activation through training under this technique. It is produced during the forward pass. It is not kept in memory. It is produced again only when the gradient step requires it. The diagram below shows that order of events. It is a simplified view of the sequence described in the source, not the internal scheduling details of any particular framework.
Notice that the discarded activation is not retrieved from storage. It is rebuilt at the moment the gradient step needs it. That rebuilding is exactly the extra work that makes training slower.
Four Layers, Two Strategies
Counting activations in memory at the start of the gradient calculation
A small hypothetical network has four layers, and each layer produces one activation during the forward pass. How many activations must sit in memory when the gradient calculation begins under two strategies: storing every activation, and using checkpointing that keeps only the activations at layers 2 and 4?
- 1Strategy one: store everything
Every layer's activation is kept for reuse, so all four activations sit in memory when the gradient calculation begins.
- 2Strategy two: keep checkpoints only
Under this illustrative choice, the activations at layers 2 and 4 are kept as checkpoints. The activations at layers 1 and 3 are discarded during the forward pass.
- 3Recompute when needed
When the gradient step needs the activation at layer 1 or layer 3, that activation is produced again from the nearby stored checkpoint instead of being read from memory.
- 4Compare the costs
Strategy two holds two activations instead of four, so memory is reduced. The cost is the extra work of producing the two discarded activations again.
Storing everything keeps 4 activations in memory. Checkpointing this way keeps 2, at the cost of recomputing 2 activations during the gradient step. The choice of which activations to keep is illustrative here; the source describes the general approach, not a fixed placement rule.
What do you think happens?
In the four-layer example with every activation stored, how many activations must sit in memory when the gradient calculation begins?
Reveal answer
4
Each of the four layers produces one activation, and storing everything keeps all four for reuse.
Paying for Memory with Time
Checkpointing turns one cost into another. Turning it on reduces the memory needed for activations, and it adds compute because discarded activations must be recomputed. These two effects move in opposite directions: more memory saved generally means more recomputation performed.
Common Mistakes When Reasoning About Checkpointing
Common mistake Assuming the weights are always the largest memory cost
Why it failsThe source states that activation memory can be much larger than the memory used by the weights, so the weights are not always the main cost.
Do this insteadEstimate both the weights and the activations before deciding where memory is going.
Common mistake Believing checkpointing changes the model
Why it failsCheckpointing is a technique for what gets stored during training. It does not alter the model's architecture or weights.
Do this insteadThink of checkpointing as a change in storage and recomputation, not in the model.
Common mistake Describing checkpointing as free memory savings
Why it failsThe source says checkpointing reduces memory requirements but increases training time because of the extra recomputation.
Do this insteadAlways state both sides: memory is reduced and training time is increased.
Common mistake Treating gradient checkpointing and activation recomputation as two different techniques
Practice: Reason About a Memory Constraint
A team is training a large model. Memory is the constraint that currently prevents the model from being trained at all. Explain whether gradient checkpointing is a reasonable choice, naming the memory effect and the time effect in your answer.
Hints
- Recall that checkpointing reduces memory requirements.
- Recall that checkpointing increases training time because of recomputation.
- Ask which constraint is tighter: memory or time.
Key Takeaways
- Activation memory can be much larger than the memory used by model weights, so the weights are not always the main memory cost.
- Weights have a fixed footprint set by the architecture, while stored activations accumulate with the work done on each batch as the network moves forward.
- Gradient checkpointing, also called activation recomputation, avoids storing activations for reuse and recomputes them when needed.
- Checkpointing is a deliberate trade: it reduces memory requirements but increases training time because of the extra recomputation.