Activation Checkpointing › Interaction with Distributed Training

Activation checkpointing does not operate in isolation: it interacts with every other component of the training infrastructure stack. Understanding these interactions is essential for configuring large-scale training runs correctly.

In pipeline-parallel training, the model is split across multiple GPUs (stages), with each stage holding a portion of the layers. Data flows through stages in sequence, and multiple micro-batches are in flight simultaneously to keep all stages busy. This creates a critical interaction with activation memory: each stage must hold activations for all micro-batches currently buffered in its pipeline.

If a stage has $m$ micro-batches in flight and each micro-batch has activation memory $M$, the total activation memory for that stage is $m \times M$. For large pipeline schedules (many micro-batches, many stages), this can be enormous. Activation checkpointing within each stage reduces the per-micro-batch activation memory, but the multiplier $m$ still applies to the checkpointed memory.

The interleaved pipeline schedule (as used in Megatron-LM) assigns each stage multiple non-contiguous chunks of layers, which reduces the number of micro-batches needed to fill the pipeline and consequently reduces the activation memory multiplier. Combining the interleaved schedule with per-layer checkpointing provides both benefits.

In tensor-parallel training, each layer's weight matrices are split across multiple GPUs along a dimension (typically the hidden dimension). When using tensor parallelism, the activation tensors are also split: each GPU holds a shard of shape $[B, T, d/\text{tp}]$ where $\text{tp}$ is the tensor parallel size. This directly reduces activation memory by a factor of $\text{tp}$, acting as another form of memory reduction that stacks multiplicatively with activation checkpointing.

When combining tensor and pipeline parallelism with activation checkpointing (the "3D parallelism" stack), the memory savings multiply. A checkpoint boundary tensor that would occupy 67 MB in the baseline might occupy only 4 MB per GPU in a $\text{tp} = 8$, $\text{pp} = 2$ configuration. Checkpointing every layer in this setting saves memory from a much smaller base, meaning you can often reduce the checkpoint frequency (fewer segments, less compute overhead) while still staying within the memory budget.

Gradient accumulation splits a large logical batch into multiple smaller micro-batches, runs forward and backward passes on each separately, accumulates the gradients, and only updates the optimizer once per logical batch. This is used to simulate large batch sizes when memory would not allow fitting the full batch in one pass.

The interaction with activation checkpointing is additive rather than multiplicative in the memory dimension: each micro-batch's activations are independent. You run forward-backward for micro-batch 1 (using checkpointed activations), then run forward-backward for micro-batch 2, and so on. Because micro-batches do not coexist in memory during the backward pass, gradient accumulation does not increase activation memory. Activation checkpointing reduces the memory within each micro-batch's pass, and gradient accumulation reduces the effective batch size per pass. Together, they allow training at large logical batch sizes with very limited per-pass memory.

Data-parallel training replicates the model on multiple GPUs and splits the batch across them. Each GPU runs a complete forward-backward pass on its local batch shard. Gradient checkpointing applies independently per GPU, since each GPU's computation is identical but on different data. The memory savings are the same as in the single-GPU case. The only interaction is that data-parallel gradient synchronization (all-reduce) happens on the gradient tensors after the backward pass, which is after all activations have been freed.