Gradient Checkpointing Trade-Offs: When Memory Savings Justify the Recompute Cost for Large Model Training

During neural network training, intermediate activation tensors from the forward pass must be retained until the backward pass computes their gradients. For transformer models, these activations consume 5-15x more GPU memory than the model parameters themselves. A Llama 3.1 405B model trained with micro-batch size 1 and sequence length 8K requires approximately 810 GB of activation memory per GPU for a single transformer layer's activations, even before accounting for optimizer states and parameters.

Gradient checkpointing (also called activation checkpointing or rematerialization) trades compute for memory by discarding selected intermediate activations during the forward pass and recomputing them on demand during the backward pass. The technique reduces activation memory by 40-80% while adding 10-30% compute overhead, depending on the checkpointing strategy and model architecture. The core question for practitioners: which activations to checkpoint and which to retain.

Full checkpointing discards all intermediate activations and recomputes every layer from the nearest saved checkpoint during backpropagation. This maximizes memory savings (up to 80% reduction) but adds approximately 33% compute overhead for standard transformer architectures, since every forward computation is effectively performed twice. Memory drops to storing only the input activations at each checkpoint boundary, typically one per transformer block.

Selective checkpointing, introduced by Chen et al. (2024), retains attention logits and softmax outputs while checkpointing only the MLP activations. Since attention logits are cheap to store (batch size x sequence length x number of heads) but expensive to recompute (requiring full QKV projections), selective checkpointing reduces memory by 50-60% with only 12-18% compute overhead. The table below compares strategies for a 70B transformer training run.

The recompute cost varies significantly by transformer sublayer. Self-attention is the most expensive to recompute because it requires re-running the full QKV projection (3 weight matrix multiplications) plus the attention score computation and output projection. MLP layers are approximately 40% cheaper to recompute per activation byte because they involve simpler linear projections without the softmax and masking operations.

On H100 SXM, recomputing one self-attention layer's activations for a 70B model requires approximately 280 TFLOPs and takes 1.7 ms per micro-batch. By contrast, recomputing the MLP activations for the same layer requires 170 TFLOPs and 1.0 ms. This asymmetry makes selective checkpointing that retains attention activations while checkpointing MLP activations the optimal trade-off for transformer architectures.

Checkpointing strategy determines the feasible model size and micro-batch size for a given GPU memory budget. On H200's 141 GB HBM3e, a Llama 3.1 70B model (140 GB weights in FP16 + 28 GB optimizer states + 12 GB gradients) leaves approximately 0 GB for activations without checkpointing. Full checkpointing at the granularity of every transformer block frees roughly 180 GB of activation memory, enabling micro-batch size 8 at sequence length 8K.

On B300's 288 GB HBM3e, the same model's weights + optimizer + gradients consume 180 GB, leaving 108 GB for activations. Without checkpointing, this limits micro-batch to approximately 4 at 8K sequence length. Selective checkpointing more than doubles the feasible micro-batch to 10-12, improving throughput by 40-50% at the cost of 15-18% recompute overhead. The net effect is a 20-30% training speed improvement over no-checkpointing with smaller batch sizes.

Modern frameworks support automatic checkpoint placement using dynamic programming to optimize the set of saved activations. PyTorch's automatic activation checkpointing (torch.distributed.checkpoint and torch.compile integration) analyzes the compute graph to identify optimal save points given a user-specified memory budget. The solver considers per-layer recompute cost, memory footprint, and the dependency graph to produce a near-optimal checkpointing schedule.

DeepSpeed's activation offloading takes a different approach: instead of recomputing, it offloads activations to CPU memory during the forward pass and prefetches them back during backward pass. CPU offloading has lower compute overhead (5-10%) than recompute checkpointing but requires high-bandwidth CPU-GPU interconnect (PCIe Gen5 or NVLink-C2C) and is limited by PCIe bandwidth of approximately 64 GB/s on standard servers. For models where CPU offloading is feasible, it achieves 90% of the memory savings of full checkpointing with only 5-10% throughput degradation.

Different training frameworks implement checkpointing with varying granularity and integration with parallelism strategies. PyTorch FSDP integrates checkpointing at the transformer block level, allowing users to wrap nn.Transformer blocks with checkpoint_wrapper for automatic forward recomputation. Megatron-LM offers tensor-parallel-aware selective checkpointing that avoids recomputing activations that are already partitioned across tensor-parallel ranks.

The key difference is memory vs throughput trade-off calibration. FSDP with activation checkpointing enables training 70B models on H200 with micro-batch size 4 at approximately 65% of the throughput achievable without checkpointing. Megatron-LM's selective checkpointing achieves approximately 80% of baseline throughput at similar memory savings, but requires manual annotation of which sublayers to checkpoint. For most teams, FSDP's automatic approach is the pragmatic choice despite slightly lower throughput.

Use selective checkpointing (checkpoint MLP activations, retain attention activations) as your default for all transformer training runs on H100-class GPUs. It delivers 50-60% memory reduction with only 12-18% compute overhead, the best memory-to-performance trade-off for transformers. Full checkpointing is only necessary when training models that barely fit into GPU memory, such as 405B+ parameter models on H200 with micro-batch size 1.

For B200 and B300 GPUs with 288 GB HBM, the larger memory budget reduces the need for aggressive checkpointing. On these GPUs, selective checkpointing with micro-batch sizes of 8-12 provides the best throughput for most models up to 300B parameters. Only use checkpointing at all when it enables a larger micro-batch that increases model flops utilization. Running without checkpointing with micro-batch size 1 is often slower than using checkpointing to enable micro-batch size 4, even accounting for the recompute overhead.