Gradient Checkpointing

Gradient checkpointing is a memory-saving training technique that stores only selected intermediate activations from the forward pass and recomputes the others during backpropagation. Also called activation checkpointing, it exchanges additional computation for lower peak memory use. Despite its name, the technique checkpoints activations rather than parameter gradients or model files, making it especially valuable when activation tensors prevent a neural network from fitting in available GPU memory.

How Gradient Checkpointing Works#

During a standard forward pass, a neural network calculates intermediate tensors called activations. The automatic differentiation system retains the activations needed to calculate gradients later, as described in the PyTorch autograd mechanics. Deep networks, large batches, high-resolution images, and long input sequences can make these saved tensors consume substantial memory.

Gradient checkpointing divides the network into segments:

1. The forward pass stores inputs or boundary activations for selected segments.
2. Other intermediate activations inside those segments are discarded.
3. During the backward pass, each checkpointed segment runs forward again to reconstruct the missing values.
4. The reconstructed activations are immediately used to calculate gradients.

The PyTorch activation checkpointing API exposes this behavior through torch.utils.checkpoint. Equivalent concepts appear as JAX gradient checkpointing and rematerialization and TensorFlow tape checkpointing.

Checkpoint placement determines the tradeoff. Checkpointing more regions generally saves more memory but repeats more operations. Selective placement around activation-heavy blocks can provide a better balance than recomputing the entire network.

Tradeoffs and Related Techniques#

Gradient checkpointing usually leaves model parameters, gradients, and optimizer states unchanged. Its primary target is activation memory, and training becomes slower because some forward computations run twice. The exact result depends on the architecture, checkpoint boundaries, batch shape, and hardware.

It differs from several related techniques:

• Gradient accumulation processes multiple microbatches before updating model weights, creating a larger effective batch without loading every sample simultaneously. Gradient checkpointing instead reduces the activations retained for each microbatch.
• Mixed precision uses lower-precision data types for selected operations. The PyTorch automatic mixed precision workflow can reduce memory and accelerate compatible operations, while checkpointing deliberately adds computation.
• Batch size reduction lowers memory by processing fewer samples together. Checkpointing may allow a larger batch or input resolution to remain feasible.
• Training checkpoints save weights and optimizer state for recovery or later inference. The PyTorch model checkpoint guide describes this persistence mechanism, which is unrelated to activation recomputation.

Checkpointed code should be functionally consistent between its original and recomputed forward executions. Mutable state, device transfers, or uncontrolled randomness inside a checkpointed region can cause errors or incorrect gradients. In decoder-style transformers, a key-value cache may also need to be disabled during checkpointed training because cached inference state can conflict with rebuilding the forward graph.

PyTorch Example#

The following example checkpointes a memory-intensive block during one training step:

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

Only the block’s inputs and required boundary information are retained; its internal activations are reconstructed during loss.backward(). Real projects should compare runs with and without checkpointing using a measurement such as PyTorch peak allocated GPU memory.

Реальные приложения#

• High-resolution computer vision: Medical segmentation, aerial detection, and industrial inspection may train on large images whose feature maps consume more memory than model weights. Checkpointing selected backbone stages can preserve image resolution or permit additional samples per batch instead of aggressively downscaling inputs.

• Long-sequence transformer training: Activation storage grows rapidly as sequence length and layer count increase. Recomputing transformer blocks can make longer contexts or larger microbatches fit on the same accelerator. The NVIDIA activation recomputation guide illustrates full and selective recomputation for transformer layers.

Practical Guidance#

Use gradient checkpointing when profiling shows that activations, rather than parameters or optimizer state, dominate memory. Start with large repeated blocks and benchmark peak memory, iteration time, and validation behavior before expanding coverage.

For documented Ultralytics YOLO training configurations, first-line memory controls include automatic mixed precision, image size, and physical batch size. The Ultralytics AutoBatch reference explains automatic batch selection based on available GPU memory. When local hardware remains insufficient, Ultralytics Platform cloud training provides configurable cloud GPUs for managed training runs. Gradient checkpointing can complement these controls when a custom PyTorch architecture requires finer activation-memory management.