Calculating Memory Requirements for Model Weights
Activation memory can be much larger than the memory used by model weights, so the weights are not always the main memory cost.
Why Weights Are Not the Whole Story
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 explains why that happens and introduces one of the main techniques used to reduce it. As you read, keep one question in mind: 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?
Where Activation Memory Builds Up
During the forward pass, the network moves an input through its layers, and each layer produces an activation. If those activations are stored so that the gradient computation can use them later, they stay in memory. The diagram below shows this for a four-layer network. Each layer adds its activation to the stored set, so the amount held in memory grows as the forward pass moves deeper.
The picture explains the growth. The weights have a fixed footprint set by the architecture, so they do not change as training runs. The stored activations, by contrast, depend on the work the network does on each batch as it moves forward. Because each stage of the forward pass contributes its own activation, the total accumulates instead of staying fixed.
Why Activations Can Outgrow the Weights
An earlier estimate in the same material assumed that the memory needed for activations was smaller than the memory needed for the model's weights. That assumption simplifies the calculation, but it does not hold for large models. The source states the point directly: activation memory can be much larger than the memory used by the weights.
The source refers to a chart comparing activation memory with weight memory across several Megatron models at different scales. The numbers come from work by Korthikanti and colleagues (2022). The main lesson is visible across those scales: as models grow, the activation side of the comparison can become the larger share. The exact values belong to the original paper and should be read there, not reconstructed from memory.
| Property | Model weights | Stored activations |
|---|---|---|
| What they are | The learned parameters of the network | Values produced during the forward pass that are kept for later gradient computation |
| What sets their size | A fixed footprint set by the architecture | The work the network does on each batch as it moves forward |
| How they change during training | They do not accumulate as the forward pass moves deeper | They accumulate, so their combined size can grow to dwarf the weights |
Weights and stored activations behave differently in memory
Storage Swapped for Computation
Gradient checkpointing, also called activation recomputation, reduces memory by not storing activations for reuse and recomputing them when they are needed.
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. The two names, gradient checkpointing and activation recomputation, 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 diagram traces one activation through this process: it is produced in the forward pass, it is not kept, and it is produced again only when the gradient step requires it. The diagram is a simplified view of the order of events, not the scheduling details of any particular framework.
Tracing Four Layers Under Two Strategies
Activations held 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 standard training and under checkpointing? The source does not prescribe which activations a checkpointing scheme keeps, so this example uses a simple illustrative choice: layers 1 and 3 are kept, and layers 2 and 4 are discarded.
- 1Standard training
Every activation is kept for reuse, so all four activations from layers 1 to 4 sit in memory when the gradient calculation begins.
- 2Checkpointing forward pass
During the forward pass, the activations from layers 1 and 3 are kept, while the activations from layers 2 and 4 are not kept for reuse.
- 3Count held in memory
Only two activations are held in memory when the gradient calculation begins, instead of four.
- 4Recomputation
When the gradient step needs the activations from layers 2 and 4, they are produced again. This is the extra work that buys the memory saving.
Standard training holds four activations at once, and this checkpointing choice holds two. The saving in memory is paid for by recomputing the two discarded activations later.
The Time Cost Behind the Memory Savings
Gradient checkpointing reduces memory requirements but increases training time because of the extra recomputation. The source names both effects and does not claim that one always beats the other. The comparison below sets the two strategies side by side.
The choice depends on which constraint is binding. When memory is the constraint that prevents a model from being trained at all, the time cost may be a reasonable price. When time is the tighter constraint and memory already fits, storing activations as usual may be the better choice. Reasoning about memory should therefore weigh both the weights and the activations, and treat checkpointing as a deliberate trade of memory for time.
Common Misreadings of Checkpointing
Common mistake Assuming the weights are always the largest memory cost
Why it failsThe source states that the weights are not always the main memory cost. For large models, activation memory can be much larger than the memory used by the weights.
Do this insteadEstimate both the weights and the stored activations before deciding where the memory goes.
Common mistake Thinking checkpointing changes the model
Why it failsCheckpointing is a technique for handling activations. It reduces memory by not storing activations for reuse, without changing the model itself.
Do this insteadTreat checkpointing as a change to how activations are stored and recomputed, not as a change to the architecture or weights.
Common mistake Treating checkpointing as free
Why it failsCheckpointing reduces memory requirements but increases training time because of the extra recomputation.
Do this insteadAlways weigh the memory saved against the added training time before choosing it.
Checking Your Understanding
Using the four-layer example from earlier, suppose a checkpointing scheme keeps the activations for layers 1 and 3 and discards the activations for layers 2 and 4. How many activations are discarded, and how many of those must be produced again during the gradient calculation? Explain why the saving in memory comes with a cost in training time.
Hints
- Count the activations that are not kept for reuse.
- Recomputation is needed for any activation that was discarded and is required by the gradient step.
- Think about what happens to the total work done in a single training step when activations are produced twice.
Key Takeaways
- For large models, activation memory can be much larger than the memory used by the weights, so the weights are not always the main memory cost.
- Activations accumulate during the forward pass because each layer produces one and they are stored for later gradient computation, while the weights keep a fixed footprint set by the architecture.
- Gradient checkpointing, also called activation recomputation, avoids storing activations for reuse and recomputes them when needed.
- Checkpointing reduces memory requirements but increases training time because of the extra recomputation, so it is a deliberate trade of memory for time.