Memory Management: Activations, Gradients

Every large language model training run is ultimately constrained by one resource before compute, before data, even before time: GPU memory. A model that fits in memory trains; one that does not, crashes. Understanding exactly what occupies GPU memory, how much each component costs, and how to diagnose failures when things go wrong separates practitioners who can scale models from those who hit walls.

This chapter breaks down the four major memory consumers in transformer training: model parameters, gradients, optimizer states, and activations. For each category, you will learn how much memory it requires, why that amount is unavoidable, and what tradeoffs appear when you try to reduce it. You will then build a working memory estimator that predicts whether a model fits in a given GPU budget before you ever start training, and finally work through systematic OOM (out-of-memory) debugging strategies to recover from crashes when estimation alone is not enough.

The concepts here connect directly to the distributed training strategies covered in the next chapters. Tensor parallelism, pipeline parallelism, and gradient checkpointing all exist primarily to manage memory. Without understanding what they are managing, you cannot reason about when to use them or what tradeoffs they impose.

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:

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$:

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:

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:

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:

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:

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

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:

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$:

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.

Estimating total memory before running a training job lets you choose the right GPU configuration, set batch sizes appropriately, and avoid wasting hours waiting for an OOM crash that was predictable from the start. A good estimate also helps you make principled decisions about which memory reduction techniques to apply and in what order.

For a standard mixed-precision training setup with Adam, the total memory $M$ breaks down as:

where $M_{\text{other}}$ covers CUDA context, framework overhead, and temporary buffers. The CUDA runtime itself requires roughly 300-500 MB when initialized. PyTorch maintains a memory caching allocator that retains freed memory for reuse; its overhead depends on the number and sizes of allocations. Temporary workspace for cuBLAS and cuDNN operations can require several hundred megabytes more. In practice, $M_{\text{other}}$ typically adds 500 MB to 2 GB depending on the framework version and GPU model.

The activation memory $M_{\text{act}}$ depends heavily on whether gradient checkpointing is used:

where the factor of 2 bytes per element (FP16) is incorporated into the constants 34 and 5.

Consider a GPT-2 style model with the following architecture:

• Vocabulary: $V = 50{,}257$
• Hidden dim: $d = 768$
• Layers: $L = 12$
• Heads: $h = 12$
• FF dim: $d_{ff} = 3{,}072$

The parameter count is approximately:

In mixed precision with Adam:

For a batch size of 8 and sequence length 512:

Total estimated memory: roughly 3.7 GB, well within a single 16 GB GPU's budget. Scaling to a 7B model with batch size 16 and sequence length 2048 pushes this into the multi-hundred-gigabyte range, requiring distributed training.

One additional nuance in the estimation is memory fragmentation. GPU memory is allocated and freed in blocks of varying sizes. Over time, freed blocks of one size may not be reusable for allocations of a different size, causing apparent "wasted" memory that shows as free but is not contiguous enough for large allocations. PyTorch's CUDA memory allocator tries to mitigate this through block pooling and splitting, but fragmentation can cause an OOM even when the total free memory nominally exceeds the requested allocation. In practice, you should maintain a 10-20% safety margin in your memory estimates to account for fragmentation.

Let us implement a memory estimator that computes these values programmatically and validates them against PyTorch's reported memory usage. The goal is a tool you can run before any training job to predict GPU requirements and identify which component is the bottleneck.

We start by defining the model configuration and implementing functions to count parameters and estimate memory:

This gives us the parameter count for our toy model. The parameter range in the tens of millions is appropriate for single-GPU experimentation, where we can validate our estimation formulas against actual PyTorch measurements without requiring expensive hardware.

Now we implement the full memory estimator using the formulas derived above:

The table shows how the optimizer state is the dominant consumer for this model configuration, a pattern that holds across all transformer scales. Gradient checkpointing has no effect on the static components (parameters, gradients, optimizer states), but it dramatically reduces the activation term. Activations become comparably large relative to static memory at longer sequence lengths, where the $T^2$ attention term dominates.

Let us build an actual small model and verify that our estimates align with what PyTorch reports:

Our analytical estimate closely matches the ground truth parameter count from PyTorch. The small discrepancy comes from minor architectural differences between our counting formula and the exact layer implementations, particularly in how nn.MultiheadAttention organizes its projection weight tensors.

We can instrument actual GPU memory consumption if a CUDA device is available. Even on CPU, we can measure the parameter storage footprint:

The match between estimated and actual parameter memory confirms our formula is accurate. In a real training run, you would observe three to seven times this parameter footprint once gradients and Adam states are added.

Let us compute memory requirements across different model sizes to build intuition for how memory scales:

This table reveals why distributed training is not optional for models above a few billion parameters. A 7B model requires over 100 GB of static memory alone, far exceeding any single GPU's capacity. The 70B model requires over a terabyte of static memory, requiring many GPUs even before considering activations. Gradient checkpointing helps with activations but does nothing for the static component, which is why memory-efficient optimizers and sharding strategies are both essential tools at large scale.

The stacked comparison shows why practitioners typically combine multiple techniques rather than relying on any single one. Mixed precision alone gets you to the 14P static baseline. Adding gradient checkpointing addresses the activation term for long sequences. Memory-efficient optimizers then chip away at the dominant static component. The full combination can make a 7B model trainable on hardware that would otherwise require tensor parallelism.

Out-of-memory errors are among the most frustrating failures in deep learning, because they often manifest far into a training run, sometimes after hours of successful execution. A systematic debugging approach dramatically reduces the time to resolution.

A typical CUDA OOM error looks like:

This message tells you several things. The model tried to allocate 4.5 GB in a single operation. The GPU has 79.2 GB total capacity. PyTorch has 71.4 GB actively in use and 73.1 GB reserved (the difference between reserved and allocated is fragmentation). Only 3.3 GB of contiguous usable memory remains, which is not enough for the 4.5 GB allocation.

The allocation that triggered the error is often not the root cause. Memory may have been building up gradually across many operations, and this particular allocation is simply the one that crossed the threshold. Debugging therefore requires the full memory timeline rather than an inspection of only the final allocation.

The fragmentation number here is particularly telling: 73.1 GB reserved but only 71.4 GB allocated means 1.7 GB of fragmented memory that PyTorch's allocator is holding but cannot currently provide as a contiguous block. When you see a large gap between "reserved" and "allocated," fragmentation is likely contributing to your OOM. The solution is often to restart the process and reduce peak allocation spikes rather than just reducing total memory usage.

Step 1: Confirm available GPU memory. Before anything else, verify you know how much GPU memory your target hardware provides, and that no other processes are consuming it:

Other processes may hold GPU memory even when they appear idle. Jupyter kernels are a common culprit: if you have previously run a training experiment in the same kernel, tensors from that run may still occupy memory. Always check nvidia-smi before attributing an OOM to your model's memory requirements.

Step 2: Add memory profiling checkpoints. Insert torch.cuda.memory_summary() calls at key points to see where memory is consumed. On a CUDA device:

Step 3: Test with minimal batch size. If the model OOMs even at batch size 1, the model itself does not fit in memory. This indicates you need either mixed precision, offloading, or model parallelism. If it runs at batch size 1, you have a scaling issue and the problem is the activation memory.

Step 4: Enable gradient checkpointing. This is typically the first optimization to reach for when activations are the bottleneck:

Step 5: Reduce batch size or sequence length. Both linearly reduce activation memory. Gradient accumulation can maintain effective batch size while using smaller micro-batches.

Step 6: Use mixed precision. If training in FP32, switching to FP16 or BF16 halves parameter and gradient memory. Most modern training code defaults to BF16 but it is worth verifying.

Step 7: Check for memory leaks. Memory that is not freed properly can accumulate across iterations. Common culprits include:

• Storing loss values without detaching (losses.append(loss) instead of losses.append(loss.item())): keeps the entire computation graph alive
• Keeping references to intermediate tensors outside of training loops
• Running validation inside torch.enable_grad() context instead of torch.no_grad()
• Logging or visualization code that retains tensors

Step 8: Try CPU offloading. If your model does not fit on GPU at all, CPU offloading moves optimizer states or model weights to CPU RAM between uses. This trades compute throughput for memory capacity. Libraries like DeepSpeed ZeRO-3 implement this automatically. The penalty is communication bandwidth between CPU and GPU, which can reduce throughput by 2-5x but makes otherwise impossible configurations feasible.

Here is a reusable profiling utility that tracks memory at each step of the training loop:

The delta column shows exactly when memory spikes occur. In a real training run on GPU, you would instrument these checkpoints around the forward pass, loss calculation, and backward pass to pinpoint where unexpected memory consumption appears. The large delta after the backward pass confirms that gradient tensors are being populated, and the cleanup step shows that proper deletion and garbage collection frees this memory.

The following table summarizes the most frequent OOM patterns encountered in practice:

When you encounter a memory constraint, the choice of which technique to apply first depends on which component is the bottleneck. The decision follows a natural ordering from cheapest to most expensive in terms of code complexity and performance impact.

If static memory (parameters + gradients + optimizer states) is the issue:

Start with mixed precision if not already enabled. This halves parameter and gradient memory with minimal code change and typically no convergence impact. If static memory is still too high after mixed precision, consider switching from Adam to 8-bit Adam (a drop-in replacement) to halve optimizer state memory. For extreme memory constraints, ZeRO sharding (discussed in the distributed training chapters) shards all static memory across GPUs, so the per-device cost scales as $1/N$ of the full model.

If activation memory is the issue:

Enable gradient checkpointing. This is almost always the right first move when activation memory is the bottleneck: it reduces activation memory by a factor of $L$ at the cost of approximately 33% more compute, and requires changing only a single line of code for HuggingFace models. If checkpointing is already enabled and you are still OOM, reduce batch size and increase gradient accumulation steps to maintain the same effective batch size. For persistent issues with long sequences, FlashAttention replaces the standard attention computation with a memory-efficient kernel that never materializes the full $T \times T$ attention matrix.

If both are issues:

Apply all of the above and consider CPU offloading for the components you use least frequently. Optimizer states are excellent candidates for CPU offloading because they are only needed once per optimizer step, not on every forward or backward pass. DeepSpeed's ZeRO-Offload moves optimizer states to CPU RAM between steps, freeing the $8P$ bytes of Adam states from GPU memory while paying a CPU-GPU transfer cost at each step.

The key configuration parameters that control memory usage during transformer training are:

• batch_size: Controls activation memory linearly. Halving batch size halves activation memory but may require more gradient accumulation steps to maintain the effective batch size.
• seq_len / max_position_embeddings: Controls both the $T^2$ attention term and the linear activation terms. Doubling sequence length roughly quadruples attention activation memory and doubles all other activation memory, for an overall factor slightly less than 4x depending on the model's dimension.
• num_layers: Controls how many layers' activations must be stored simultaneously. Gradient checkpointing amortizes this to one layer regardless of depth.
• hidden_dim: Controls the linear activation memory per token and the parameter count per layer. Scaling hidden dimension has a quadratic effect on parameter count (through the $d^2$ attention projection terms) and a linear effect on per-token activation memory.
• optimizer: Adam requires 8 bytes per parameter for states; alternatives like Lion (4 bytes), SGD (4 bytes), or 8-bit Adam (2 bytes) can significantly reduce this. AdaFactor's factored representation provides sub-linear scaling but requires careful learning rate tuning.
• param_dtype: BF16 parameters halve parameter memory compared to FP32. Mixed precision typically uses FP16/BF16 parameters with FP32 gradients and optimizer states.
• gradient_checkpointing: Reduces activation memory from $O(L \times B \times T \times d)$ to $O(B \times T \times d)$ at the cost of approximately 33% more forward pass compute.

Memory estimation formulas are approximations, not guarantees. Real GPU memory usage includes framework overhead, CUDA kernel buffers, temporary tensors created during forward pass operations, and memory fragmentation that prevents large contiguous allocations even when total free memory appears sufficient. The estimator built here provides a useful floor estimate; actual usage often runs 10-20% higher due to these factors. When planning a training run, always target 80-85% GPU utilization in your memory budget rather than 100%, to leave room for fragmentation and framework overhead.

The interaction between memory and compute efficiency is also important to understand. Many memory reduction techniques trade memory for compute. Gradient checkpointing adds 33% compute overhead. CPU offloading adds PCIe bandwidth latency. Mixed precision improves both memory and compute throughput on hardware with tensor core support, making it the only "free" memory reduction. Understanding these tradeoffs lets you make rational choices: if your training is already compute-bound (GPU utilization near 100%), adding gradient checkpointing will slow training noticeably. If you are memory-bound and the GPU is sitting idle between steps, the checkpointing compute overhead costs you nothing.

Memory management in practice also interacts heavily with the distributed training strategies covered in subsequent chapters. Tensor parallelism splits individual weight matrices across GPUs, reducing per-device parameter memory but requiring high-bandwidth communication for every matrix multiplication. Pipeline parallelism assigns different layers to different GPUs, reducing parameter memory per device but requiring careful management of activation buffers at pipeline stage boundaries. Fully Sharded Data Parallelism (FSDP) shards parameters, gradients, and optimizer states across all GPUs, making the effective per-device memory proportional to $1/N$ of the full model, where $N$ is the GPU count.

The "14x rule" for static memory deserves a practical caveat: it assumes the full training stack is running. Fine-tuning with LoRA or other parameter-efficient methods dramatically changes the breakdown. In LoRA fine-tuning, only the adapter parameters receive gradients and optimizer states, while the base model parameters are frozen. For a 7B model fine-tuned with LoRA at rank 16, the adapter might have fewer than 10M trainable parameters out of 7B total. The gradient and optimizer state costs apply only to those 10M parameters, while the base model parameters require only inference memory ($2P$ bytes for FP16). This can reduce the effective multiplier from 14x to something closer to 2-3x of the base model size, making fine-tuning of large models feasible on single GPUs that could not possibly run full training.

Understanding these tradeoffs is not merely academic. Production training runs at scale involve budgeting across hardware configurations, choosing between larger batches with gradient checkpointing versus smaller batches without it, and deciding which distributed strategy minimizes communication overhead while keeping per-device memory within budget. The memory breakdown covered in this chapter is the foundation for all of those decisions.

Memory management in transformer training centers on four categories of consumption, each with distinct characteristics:

• Parameters cost $2P$ bytes in FP16, setting a hard floor on GPU memory requirements. They scale with the architecture's depth and width, and they persist for the entire training run.
• Gradients cost $4P$ bytes in FP32 for numerical stability, matching the parameter shape exactly. They exist only during backpropagation and are freed after each optimizer step.
• Optimizer states cost $8P$ bytes for Adam (two FP32 moments), making the optimizer the single largest memory consumer. Alternatives like Lion (4 bytes), 8-bit Adam (2 bytes), and AdaFactor (sub-linear) reduce this significantly at varying convergence tradeoffs.
• Activations cost $O(L \times B \times T \times d)$ bytes, growing quadratically with sequence length due to attention weight matrices. Gradient checkpointing reduces this to $O(B \times T \times d)$ at the cost of approximately 33% more compute.

The combined static memory for mixed-precision Adam training follows the "14x rule": $14P$ bytes. For a 7B model, this is approximately 98 GB before activations. Memory estimation before starting a training run prevents wasted hours on predictable OOM crashes. When crashes do occur, a systematic approach of profiling checkpoints, testing at minimal batch size, and progressively applying reduction techniques provides a clear path to recovery.

When you encounter memory pressure, apply techniques in order of their cost-effectiveness: mixed precision first (free performance improvement), gradient checkpointing second (33% compute overhead), memory-efficient optimizers third (no compute overhead but may affect convergence), and distributed sharding or CPU offloading last (high complexity but unlimited scalability). Combining all three of the first techniques typically makes 7B-class models trainable on a single A100 80GB GPU, which represents the boundary between single-GPU and multi-GPU training for the current generation of hardware.

The next chapter extends these memory concepts into data parallelism, where replicating the model across GPUs multiplies parameter and optimizer state copies while enabling larger effective batch sizes.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about memory management in transformer training.