Memory Management › The Four Memory Consumers

GPU memory during training is divided among four distinct categories, each with different characteristics and different strategies for reduction. The total memory budget must accommodate all four simultaneously, and failing to account for any one of them is a common source of mysterious OOM errors late in training runs.

One useful framing: think of parameter memory, gradient memory, and optimizer state memory as static costs that depend only on model architecture. They exist as soon as the model and optimizer are instantiated and they persist until training ends. Activation memory, by contrast, is a dynamic cost that depends on the current batch: its size, the sequence length, and whether you have enabled techniques to reduce it. The two categories require different strategies for reduction and different tools for debugging.

Model parameters are the learned weights stored in the model: embedding matrices, attention projection matrices, feed-forward network weights, layer normalization scales and biases, and output projection weights. These are the values that exist before training begins (initialized randomly or from a checkpoint) and persist throughout the entire training process.

For a transformer with $P$ parameters, the parameter memory in bytes depends directly on the numeric precision used:

• In full 32-bit float (FP32): $4P$ bytes
• In 16-bit half-precision float (FP16 or BF16): $2P$ bytes
• In 8-bit integer (INT8): $P$ bytes

Modern large-scale training almost universally uses BF16 or FP16 for parameters during the forward pass, so a model with 7 billion parameters requires roughly $2 \times 7 \times 10^9 = 14$ GB just for parameters. A 70B model requires 140 GB. These numbers set a hard floor: you need at least this much memory before accounting for anything else.

The count of parameters $P$ in a transformer scales predictably. For a model with:

• Vocabulary size $V$
• Hidden dimension $d$
• Number of layers $L$
• Feedforward intermediate dimension $d_{ff}$ (typically $4d$)
• Number of attention heads $h$, with head dimension $d_h = d/h$

The dominant parameter contributions are:

$$P \approx V \cdot d + L \cdot \left(4d^2 + 2 \cdot d \cdot d_{ff}\right) + d$$

where the first term covers the embedding table, the second covers attention projections and feed-forward weights per layer, and the third covers the final layer norm. The attention projection contribution of $4d^2$ comes from four projection matrices $W_Q$, $W_K$, $W_V$, $W_O$, each of shape $d \times d$. The feedforward contribution $2 \cdot d \cdot d_{ff}$ comes from two linear layers: one expanding from $d$ to $d_{ff}$ and one projecting back to $d$.

To make this concrete, for a 7B-parameter model with $d = 4096$, $L = 32$, $d_{ff} = 11{,}008$, $V = 32{,}000$:

$$\begin{aligned} \text{Embeddings} &= 32{,}000 \times 4096 \approx 131M \\ \text{Per-layer attention} &= 4 \times 4096^2 \approx 67M \\ \text{Per-layer FFN} &= 2 \times 4096 \times 11{,}008 \approx 90M \\ \text{Per-layer total} &\approx 157M \\ \text{All 32 layers} &\approx 5{,}024M \\ \text{Grand total} &\approx 6{,}700M \approx 7B \end{aligned}$$

This derivation shows how the bulk of parameters live in the transformer layers themselves. The embedding table is a significant contributor for models with large vocabularies, but for modern decoder-only models the layer weights dominate.

During backpropagation, the training loop computes the gradient of the loss with respect to every parameter. These gradients have exactly the same shape as the parameters themselves, meaning they require an identical amount of memory. If you store parameters in FP16 ($2P$ bytes), you need additional memory for the corresponding gradients.

The precision question for gradients is subtler than for parameters. In mixed-precision training, the forward pass runs in FP16 or BF16 to take advantage of fast half-precision matrix multiplication hardware. But accumulating small gradient signals across millions of samples in FP16 is numerically dangerous: the limited mantissa precision causes small gradients to underflow to zero and large ones to overflow. The standard solution is to store and accumulate gradients in FP32, which provides enough dynamic range and precision to keep the update signal intact. This means gradients often cost $4P$ bytes even when parameters themselves are stored in FP16.

In practice, the mixed-precision pipeline looks like this: during the forward pass, parameters are cast from their FP32 master copies to FP16 for the actual computation. The resulting FP16 activations flow through the network. During the backward pass, gradients are computed in FP16 and then immediately cast to FP32 before accumulation. The optimizer update then operates entirely in FP32. After the update, the updated FP32 parameters are cast back to FP16 for the next forward pass. The FP32 master copy exists to accumulate small updates that would otherwise vanish in FP16.

Gradients exist only during backpropagation. They do not accumulate across batches (unless you are using gradient accumulation, in which case you are intentionally summing them). After the optimizer applies the update and zero_grad() is called, this memory is freed. However, at peak training memory usage, the gradient tensors are fully populated and occupy their full allocation.

Gradient memory can be reduced through gradient compression (sending lower-precision gradients in distributed settings), gradient accumulation (which does not reduce peak memory but allows using smaller batches), and gradient checkpointing (which trades compute for activation memory, not directly for gradient memory). The gradient footprint itself is fundamentally tied to parameter count.

The optimizer is often the largest single memory consumer, exceeding both parameters and gradients. This surprises many practitioners who focus on model size without accounting for the optimizer.

Adam, the most commonly used optimizer for transformer training, maintains two moving average tensors for every parameter:

• The first moment estimate $m_t$ (exponentially decaying average of gradients)
• The second moment estimate $v_t$ (exponentially decaying average of squared gradients)

The Adam update rule ties these together. Given the gradient $g_t$ at step $t$, Adam computes:

$$\begin{aligned} m_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ v_t &= \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \\ \hat{m}_t &= \frac{m_t}{1 - \beta_1^t} \quad \text{(bias correction)} \\ \hat{v}_t &= \frac{v_t}{1 - \beta_2^t} \quad \text{(bias correction)} \\ \theta_t &= \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t \end{aligned}$$

where:

• $\beta_1 \approx 0.9$ is the decay rate for the first moment
• $\beta_2 \approx 0.999$ is the decay rate for the second moment
• $\eta$ is the learning rate
• $\epsilon$ is a small constant for numerical stability (typically $10^{-8}$)

Both $m_t$ and $v_t$ are typically stored in FP32, regardless of parameter precision. For a model with $P$ parameters, Adam requires:

$$\text{Adam optimizer memory} = 4P + 4P = 8P \text{ bytes}$$

This means that for a 7B parameter model, Adam alone requires approximately $8 \times 7 \times 10^9 = 56$ GB. Combined with FP16 parameters ($14$ GB) and FP32 gradients ($28$ GB), the total before activations is $14 + 28 + 56 = 98$ GB for a 7B model trained with Adam in mixed precision.

The reason both moments are stored in FP32 rather than FP16 is numerical stability. The second moment estimate $v_t$ accumulates squared gradient values, and the Adam update divides by $\sqrt{v_t}$. Very small values of $v_t$ in low-precision arithmetic lead to division instability and divergence. The first moment $m_t$ is similarly vulnerable: if the gradient signal is small and $m_t$ underflows to zero in FP16, the optimizer loses the ability to update parameters with slowly changing gradients. FP32 provides enough dynamic range to prevent both failure modes.

Alternative OptimizersLink Copied

When Adam's $8P$ bytes is too expensive, several alternatives offer lower memory overhead with varying tradeoffs:

Adafactor approximates the second moment using a factored representation. For a weight matrix $W$ of shape $m \times n$, Adam stores a full $m \times n$ second moment matrix. Adafactor instead stores a single row vector of size $m$ and a column vector of size $n$, whose outer product approximates the full matrix. This reduces memory from $O(mn)$ to $O(m + n)$ for each matrix. Across the model, optimizer memory drops from $O(P)$ to roughly $O(\sqrt{P})$. Adafactor was used to train T5 and many subsequent encoder-decoder models and works well when the effective learning rate schedule is managed carefully.

Lion (Evolved Sign Momentum) uses only the sign of the gradient update, maintaining a single exponential moving average of gradients rather than two. The update rule is:

$$\begin{aligned} c_t &= \text{sign}(\beta_1 m_{t-1} + (1 - \beta_1) g_t) \\ m_t &= \beta_2 m_{t-1} + (1 - \beta_2) g_t \\ \theta_t &= \theta_{t-1} - \eta \cdot c_t \end{aligned}$$

Lion requires only one FP32 state tensor per parameter, halving optimizer memory from $8P$ to $4P$ bytes. Because it applies uniform update magnitude (scaled by learning rate), it is effectively a sign-gradient method and needs a lower learning rate than Adam to achieve comparable results.

SGD with Momentum stores only the velocity buffer ($4P$ bytes in FP32), consuming half the memory of Adam. It is rarely used for transformer pretraining because it converges more slowly and is more sensitive to learning rate choices. However, it can be effective for fine-tuning tasks where the model is already near a good solution.

8-bit Adam quantizes optimizer states to 8-bit integers using dynamic quantization. Each block of optimizer state values is stored with a per-block scaling factor. This reduces the $8P$ bytes of standard Adam to $2P$ bytes while maintaining similar training dynamics. The bitsandbytes library provides a drop-in replacement that works with most PyTorch training code.

Activations are the intermediate tensors computed during the forward pass, stored in memory so that backpropagation can compute gradients. Unlike parameters, gradients, and optimizer states, activation memory scales with model size, batch size, and sequence length. This dynamic behavior makes activation memory the hardest component to predict and the one most often responsible for late-stage OOM crashes.

What Activations Are StoredLink Copied

To understand why so much memory is needed for activations, it helps to walk through exactly what needs to be stored for a single transformer layer during the backward pass.

Consider a single transformer layer processing an input tensor $X$ of shape $(B, T, d)$, where $B$ is batch size, $T$ is sequence length, and $d$ is the hidden dimension. The forward pass computes:

1. Layer norm pre-attention: The normalized output $\hat{X}_1$ of shape $(B, T, d)$ must be retained to compute the layer norm gradient.
2. Query, Key, Value projections: The tensors $Q = \hat{X}_1 W_Q$, $K = \hat{X}_1 W_K$, $V = \hat{X}_1 W_V$, each of shape $(B, T, d)$, must be retained to compute projection weight gradients.
3. Attention score matrix: The raw scores $S = QK^\top / \sqrt{d_h}$ of shape $(B, h, T, T)$ must be retained for the softmax gradient.
4. Attention weights after softmax: The softmax output $A = \text{softmax}(S)$ of shape $(B, h, T, T)$ must be retained for the value gradient.
5. Attention output: The attended values $AV$ of shape $(B, T, d)$ must be retained.
6. Output projection input: The pre-projection activations for $W_O$.
7. Layer norm pre-FFN: The normalized input to the feedforward block.
8. FFN intermediate activations: The output of the first linear layer (before the activation function) of shape $(B, T, d_{ff})$.
9. FFN activation output: The post-GELU or post-ReLU tensor, shape $(B, T, d_{ff})$.
10. Layer norm values: The mean and variance used in layer normalization must be stored to compute layer norm gradients.

The attention weight matrices at items 3 and 4 are particularly expensive. Each of the $L$ layers stores two attention weight matrices of shape $(B, h, T, T)$. At FP16 (2 bytes), each matrix costs:

$$2 \times B \times h \times T^2 \text{ bytes}$$

For a 32-layer model with $B=4$, $h=32$, $T=2048$:

$$2 \times 4 \times 32 \times 2048^2 \approx 2.1 \text{ GB per attention matrix}$$

With two matrices (scores and weights) and 32 layers, the attention-only activation memory is approximately 134 GB. This is where the quadratic term dominates.

The Activation Memory FormulaLink Copied

Summing across all components of a single transformer layer and $L$ layers, the activation memory for standard full storage is approximately:

$$M_{\text{act}} \approx L \cdot B \cdot T \cdot \left(34d + 5 h T\right) \text{ bytes}$$

where the factor of 2 bytes per element (FP16) is absorbed into the constants. The first term $34d$ captures all the linear activations (projections, layer norms, FFN intermediate tensors), and the second term $5hT$ captures the quadratic attention weight matrices. The $T^2$ dependence in the second term means:

• Doubling sequence length quadruples the attention activation cost
• For long sequences, attention activations dominate over all other activation memory

For a 7B parameter GPT-style model with $d=4096$, $L=32$, $h=32$, $B=4$, and $T=2048$:

$$\begin{aligned} M_{\text{act}} &\approx 32 \times 4 \times 2048 \times (34 \times 4096 + 5 \times 32 \times 2048) \\ &= 32 \times 4 \times 2048 \times (139{,}264 + 327{,}680) \\ &= 32 \times 4 \times 2048 \times 466{,}944 \\ &\approx 122 \text{ GB} \end{aligned}$$

This activation footprint nearly matches the static memory cost. At longer sequence lengths or larger batch sizes, activations quickly become the dominant consumer.

Gradient CheckpointingLink Copied

Gradient checkpointing directly addresses the $O(L \times B \times T)$ activation memory problem. The key insight is that activations are needed for backpropagation but not for the forward pass itself. If you are willing to recompute them during the backward pass, you can discard them after the forward pass and only keep enough state to perform the recomputation.

In its simplest form, gradient checkpointing works as follows:

1. During the forward pass: Do not store activations for backpropagation. Instead, store only the input to each checkpointed segment (typically one transformer layer or a small group of layers).
2. During the backward pass: When the backward algorithm needs the activations of a layer, rerun the forward pass for that segment from its stored input. Use the recomputed activations to compute gradients, then discard them again.

The memory trade is: instead of storing all $L$ layers' activations simultaneously, you store only one layer's activations at any given time (the one currently being processed in the backward pass). This reduces activation memory from $O(L \times B \times T \times d)$ to approximately $O(B \times T \times d)$.

The compute cost is additional forward passes. For each layer, you run the forward computation twice: once during the original forward pass, and once during the backward pass. Across all $L$ layers, this means roughly $L + L = 2L$ forward passes instead of $L$, an overhead of approximately 33% more compute (since the backward pass itself is roughly as expensive as the forward pass, the total goes from $1F + 1B$ to $1F + L \cdot \frac{1}{L}F + 1B = 2F + 1B$, which for $F \approx B$ means about 33% more total compute).

PyTorch implements this through torch.utils.checkpoint.checkpoint(). For HuggingFace models, the convenience method model.gradient_checkpointing_enable() applies checkpointing to every transformer layer. The tradeoff is always worth it when you are activation-memory-limited: 33% more compute to save 10-100x activation memory is a favorable exchange in most training scenarios.

A more granular strategy is selective checkpointing: store activations for computationally cheap operations (like layer norms) but recompute activations for expensive operations (like attention matrices). FlashAttention, which we cover in Part 14, takes this idea further by restructuring the attention computation to never fully materialize the $T \times T$ attention weight matrix, saving both memory and bandwidth.