Memory Management › Memory Reduction Strategies: A Decision Framework

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.