Abstract

Training large neural networks is as much a memory management problem as a compute problem. A 70B parameter model in bf16 requires roughly 140 GB just for weights; gradients and optimizer states from Adam double or triple that figure before a single activation is stored. Gradient checkpointing, which trades recomputation for memory, has been standard practice since the Chen et al. paper in 2016, but the approach has been substantially refined. Combined with activation offloading to CPU or NVMe, selective recomputation policies, and recent work on fused backward passes, teams can now train models that would otherwise not fit in a given memory budget without resorting to more GPUs.

Gradient Checkpointing Fundamentals and the Recompute Trade-off

Standard backpropagation stores every intermediate activation during the forward pass so gradients can be computed during the backward pass. For a deep network with many layers, these activations dominate memory usage. Gradient checkpointing discards most activations after the forward pass and recomputes them layer-by-layer during the backward pass when they are needed. The memory cost drops from O(N) in depth to roughly O(sqrt(N)) with optimal checkpoint placement. The compute cost rises by about 33% on average, a trade-off that is favorable whenever memory is the binding constraint. PyTorch’s torch.utils.checkpoint implements this transparently for arbitrary modules, though naive application to every layer incurs more recomputation overhead than necessary.

Selective Checkpointing and Offloading

The efficiency gain comes from choosing which activations to checkpoint. Attention layers, particularly the KV activations in multi-head attention, are large but cheap to recompute. MLP activations at the bottleneck dimension are smaller but correspondingly less costly to keep. Megatron-LM and its derivative codebases implement selective activation recomputation that checkpoints full transformer blocks but retains activations for the attention softmax computation specifically, hitting a middle ground between full storage and full recomputation. Activation offloading to CPU RAM is a complementary technique: instead of discarding activations, they are streamed to host memory during the forward pass and streamed back during the backward pass. On systems with fast PCIe or NVLink connectivity to the host, this can recover the compute cost of recomputation at the price of bandwidth. Libraries like DeepSpeed ZeRO-Infinity and Colossal-AI both support CPU offloading with async prefetch to hide latency.

Practical Guidance for Practitioners

The right configuration depends on the ratio of memory headroom to compute budget. If training is memory-bound but compute-underutilized, aggressive checkpointing and offloading are cheap. If training is already compute-saturated, recomputation overhead becomes expensive. Profiling with PyTorch’s memory snapshot tooling (torch.cuda.memory._snapshot) before choosing a checkpointing strategy is worthwhile. For teams using FlashAttention-2 or -3, note that FlashAttention already performs its own recomputation internally during the backward pass, so applying gradient checkpointing to attention blocks on top of FlashAttention double-recomputes those operations and should be disabled for that subset of layers.