Activation Checkpointing and Selective Rematerialization: Mathematical Foundations, Sub-Linear Memory Scaling, Recomputation Trade-Offs, and Megatron-LM Attention Partitioning

Training modern large language models requires orchestrating hundreds of billions of parameters across distributed GPU clusters. While parameter counts and optimizer states are fixed for a given model architecture, the activation memory generated during forward passes scales directly with sequence length, batch size, and network depth. In long-context pre-training, activation memory rapidly eclipses static parameter weights, frequently triggering out-of-memory (OOM) faults on modern accelerator hardware.

Activation checkpointing (also known as gradient checkpointing or rematerialization) trades compute for memory by discarding intermediate activation tensors during the forward pass and selectively recomputing them during backpropagation. This explainer covers the mathematical foundations of sub-linear memory scaling, optimal segment partitioning, transformer-specific activation dynamics, Megatron-LM selective recomputation, and compiler-level graph partitioning.

1. The Anatomy of LLM Training Memory

During distributed training using 16-bit mixed precision (BF16 or FP16) and the standard AdamW optimizer, total GPU high-bandwidth memory (HBM) is partitioned into static states and dynamic activations.

Static Memory Footprint

For a model with $\Phi$ parameters:

• Model Parameters: 2 bytes per parameter ($2\Phi$ bytes) in BF16/FP16.
• Parameter Gradients: 2 bytes per parameter ($2\Phi$ bytes) in BF16/FP16.
• Optimizer States (AdamW): 12 bytes per parameter ($12\Phi$ bytes), consisting of 4 bytes for FP32 master weights, 4 bytes for first-order momentum ($m_t$), and 4 bytes for second-order uncentered variance ($v_t$).

Under standard Data Parallelism, the static memory footprint equals $16\Phi$ bytes. With ZeRO-3 or FSDP sharding across $N_{data}$ ranks, this reduces to $\frac{16\Phi}{N_{data}}$ bytes.

Dynamic Activation Memory

Activation memory consists of all intermediate tensors generated during the forward pass that must be retained in HBM to evaluate vector-Jacobian products (VJPs) during the backward pass. For a standard transformer layer with hidden dimension $h$, number of attention heads $a$, head dimension $d_k = h/a$, intermediate MLP dimension $d_{ffn} = \frac{8}{3}h$ (SwiGLU), sequence length $s$, and micro-batch size $b$, the uncheckpointed activation memory per layer in 16-bit precision comprises:

• Multi-Head Attention Q, K, V Projections: $3 \times 2 b s h = 6 b s h$ bytes.
• Pre-Softmax Attention Matrix ($QK^T / \sqrt{d_k}$): $2 b a s^2$ bytes.
• Post-Softmax Attention Probabilities: $2 b a s^2$ bytes.
• Attention Dropout Mask: $1 b a s^2$ bytes (or PRNG seed buffer).
• Attention Value Aggregation ($P V$): $2 b s h$ bytes.
• Attention Output Linear Projection: $2 b s h$ bytes.
• Attention LayerNorm / RMSNorm and Residual Buffer: $4 b s h$ bytes.
• SwiGLU MLP Gate and Up Projections: $2 \times 2 b s d_{ffn} = \frac{32}{3} b s h$ bytes.
• SiLU Activation Output Buffer: $2 b s d_{ffn} = \frac{16}{3} b s h$ bytes.
• MLP Down Projection: $2 b s h$ bytes.
• MLP LayerNorm / RMSNorm and Residual Buffer: $4 b s h$ bytes.

Summing these terms yields the total activation footprint per transformer layer without FlashAttention or checkpointing:

$$A_{layer} \approx 34 b s h + 5 b a s^2 \text{ bytes}$$

Across $L$ transformer layers, total activation memory scales as $\mathcal{O}(L \cdot b \cdot s \cdot h + L \cdot b \cdot a \cdot s^2)$. At long context windows (for example, $s = 32{,}768$ or $131{,}072$), activation memory completely dwarfs static model weights.

2. Mathematical Foundations of Gradient Checkpointing

The fundamental formulation of gradient checkpointing was established by Chen et al. (2016) in "Training Deep Nets with Sublinear Memory Cost".

Consider a feedforward computational graph of $N$ sequential layers:

$$x_0 \xrightarrow{f_1} x_1 \xrightarrow{f_2} x_2 \dots \xrightarrow{f_N} x_N = \mathcal{L}$$

where $x_k = f_k(x_{k-1}; \theta_k)$ and $\mathcal{L}$ is the scalar training objective.

Standard Backprop (Store All Activations):
Forward:   [x0] ---> [x1] ---> [x2] ---> [x3] ---> [x4] ---> [L]
Memory:    x0, x1, x2, x3, x4 stored simultaneously -> O(N) memory
Backward:  [dL/dx4] -> [dL/dx3] -> [dL/dx2] -> [dL/dx1] -> [dL/dx0]

Activation Checkpointing (Store Checkpoints Only):
Forward:   [x0]* ---> [x1] ---> [x2]* ---> [x3] ---> [x4]* ---> [L]
Memory:    x0, x2, x4 stored -> O(sqrt(N)) memory
Backward:  Recompute x3 from x2 -> Backprop layer 4 & 3 -> Discard x3
           Recompute x1 from x0 -> Backprop layer 2 & 1 -> Discard x1

Standard Backpropagation vs. Full Recomputation

1. Standard Backpropagation:

• Stores all activations $\{x_0, x_1, \dots, x_N\}$.
• Memory Complexity: $\mathcal{O}(N)$.
• Compute Complexity: One forward pass ($C_{fwd}$) and one backward pass ($C_{bwd} \approx 2 C_{fwd}$). Total FLOPs: $3 C_{fwd}$.

1. Naive Full Recomputation (Memory Minimal):

• Stores only the initial input $x_0$.
• To compute gradients for layer $k$, recompute forward from $x_0$ to $x_{k-1}$.
• Memory Complexity: $\mathcal{O}(1)$ (one layer buffer).
• Compute Complexity: $\sum_{k=1}^N k C_{layer} = \frac{N(N+1)}{2} C_{layer} = \mathcal{O}(N^2)$ (computationally intractable).

Sub-Linear Uniform Partitioning

Chen et al. proposed partitioning the $N$ layers into $k$ equal segments of length $m = N/k$.

During the forward pass:

• Store only the $k$ boundary states (checkpoints): $\mathcal{S} = \{x_0, x_m, x_{2m}, \dots, x_N\}$.
• Discard all intermediate activations within each segment.

During the backward pass through segment $j$ (spanning layers from $(j-1)m$ to $jm$):

• Load checkpoint $x_{(j-1)m}$.
• Recompute the $m-1$ intermediate activations within segment $j$ using forward operators $\{f_{(j-1)m+1}, \dots, f_{jm}\}$.
• Execute backward propagation through segment $j$ using the materialized activations.
• Free the intermediate activations immediately before moving to segment $j-1$.

Optimal Segment Derivation

The peak activation memory $M(k)$ is the sum of the $k$ stored checkpoints and the $m = N/k$ recomputed activations within the active segment:

$$M(k) = k + \frac{N}{k}$$

To find the minimum memory configuration, differentiate with respect to $k$ and set to zero:

$$\frac{d M(k)}{d k} = 1 - \frac{N}{k^2} = 0 \implies k^* = \sqrt{N}$$

Substituting $k^* = \sqrt{N}$ back into the memory objective yields:

$$m^* = \frac{N}{\sqrt{N}} = \sqrt{N}$$

$$M^* = \sqrt{N} + \sqrt{N} = 2\sqrt{N} = \mathcal{O}(\sqrt{N})$$

By partitioning a network into $\sqrt{N}$ segments of length $\sqrt{N}$, peak activation memory scales with the square root of network depth.

Exact FLOPs Overhead Analysis

Let $C_{fwd}$ denote the FLOPs required for a single forward pass over all $N$ layers. Each layer is executed exactly twice: once during the initial forward pass, and once during the segment recomputation pass in backward propagation.

Total computation required:

$$C_{total} = C_{fwd} + C_{fwd} + C_{bwd} = 2 C_{fwd} + 2 C_{fwd} = 4 C_{fwd}$$

Relative to standard training without recomputation ($C_{baseline} = C_{fwd} + C_{bwd} = 3 C_{fwd}$), the compute overhead is:

$$\text{Overhead} = \frac{4 C_{fwd} - 3 C_{fwd}}{3 C_{fwd}} = \frac{1}{3} \approx 33.33\%$$

Full activation checkpointing reduces activation memory by $\mathcal{O}(\sqrt{N})$ at the cost of a deterministic $33.3\%$ increase in total floating-point operations.

3. Binomial Tree Checkpointing and Multi-Level Schedules

For deeply recurrent or ultra-deep computation graphs, uniform single-level checkpointing can be generalized to multi-level recursive schedules.

Griewank and Walther (2000) introduced the Revolve algorithm, proving that for a computational graph of length $N$ with at most $K$ allowable memory slots, the optimal schedule follows a binomial tree structure.

Level 0: [x0] -----------------------------------------------------> [xN]
Level 1: [x0] --------------------> [x(N/2)] ----------------------> [xN]
Level 2: [x0] ---------> [x(N/4)] -> [x(N/2)] ----------> [x(3N/4)] -> [xN]

The maximum number of forward steps $N_{max}$ that can be evaluated using $K$ checkpoints with at most $R$ recomputations per layer satisfies the binomial coefficient:

$$N_{max}(K, R) = \binom{K + R}{K}$$

When $R$ is scaled proportionally, logarithmic memory complexity $\mathcal{O}(\log N)$ is achievable with $\mathcal{O}(N \log N)$ total computation.

4. Megatron-LM Selective Activation Recomputation

While full activation checkpointing trades $33.3\%$ additional compute for memory, not all operations inside a Transformer layer share the same memory-to-compute ratio.

In "Reducing Activation Recomputation in Large Transformer Models" (Korthikanti et al., 2022), NVIDIA researchers analyzed the operational intensity (FLOPs per byte of activation memory) across individual Transformer components.

Operational Intensity Discrepancy

Transformer operations divide sharply into two categories:

1. Matrix Multiplications (GEMMs):

• Linear projections ($W_q, W_k, W_v, W_o, W_{gate}, W_{up}, W_{down}$).
• Compute FLOPs per layer: $\approx 24 b s h^2 + 4 b s^2 h$.
• Activation Memory: $\approx 10 b s h$ elements.
• High operational intensity: millions of FLOPs per transferred byte. Recomputing GEMMs is computationally expensive.

1. Elementwise and Normalization Operations:

• LayerNorm / RMSNorm, Softmax, Dropout, GELU / SwiGLU activations.
• Compute FLOPs per layer: $\approx 10 b s h + 2 b a s^2$ (less than $3\%$ of total layer FLOPs).
• Activation Memory: Stores large intermediate tensors, including the $\mathcal{O}(s^2)$ attention score matrix and non-linear inputs.
• Low operational intensity: extremely small compute cost, but disproportionately large activation memory footprint.

Transformer Layer Operational Distribution:
+-----------------------------------+-----------------------------------+
| Operation Type                    | FLOP Share  | Activation Mem Share |
+-----------------------------------+-----------------------------------+
| GEMMs (Projections, Linear layers)| ~96 - 98%   | ~30%                |
| Non-GEMMs (Norms, Softmax, Act)   | ~2 - 4%     | ~70%                |
+-----------------------------------+-----------------------------------+

The Selective Recomputation Policy

Selective activation recomputation targets only the low-intensity, high-memory non-GEMM operations:

• Checkpointed Tensors: Retain the output activations of all linear projections and attention context projections.
• Discarded / Recomputed Tensors: Discard and recompute RMSNorm/LayerNorm inputs, Attention Softmax distributions ($P$), Attention Dropout bitmasks, and SwiGLU non-linear gating tensors.

By recomputing only non-GEMMs during the backward pass:

• Activation memory drops by over $60\%$ to $75\%$ per layer, completely removing the quadratic $\mathcal{O}(s^2)$ intermediate state.
• Compute overhead drops from $33.3\%$ down to approximately $2\%$ to $4\%$.

Comparison of Memory Scaling per Transformer Layer:
- No Checkpointing:       34 b s h + 5 b a s^2 bytes (FLOP overhead: 0%)
- Selective Checkpoint:   10 b s h bytes             (FLOP overhead: ~3%)
- Full Checkpointing:      2 b s h bytes             (FLOP overhead: ~33.3%)

5. Implementation in PyTorch: Reentrant vs. Non-Reentrant Autograd

Modern PyTorch provides native activation checkpointing via torch.utils.checkpoint. Understanding the underlying autograd mechanisms prevents subtle memory leaks and distributed synchronization bugs.

PyTorch Reentrant Checkpointing (use_reentrant=True)

In legacy PyTorch checkpointing, the recomputation pass executes inside a separate torch.autograd.backward invocation within a custom torch.autograd.Function.

1. Forward Pass: Disables autograd tracking (torch.no_grad()) inside the checkpointed block and saves only the input tensors.
2. Backward Pass: Executes a nested autograd forward pass from the saved inputs with gradients enabled, followed immediately by an internal backward pass.

Limitations:

• Breaks autograd hooks (such as register_hook) attached to intermediate tensors.
• Incompatible with PyTorch Distributed Data Parallel (DDP) find_unused_parameters logic.
• Adds Python interpreter recursion overhead.

PyTorch Non-Reentrant Checkpointing (use_reentrant=False)

Modern PyTorch (v2.0+) implements checkpointing using torch.autograd.graph.saved_tensors_hooks.

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

class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.norm1 = nn.RMSNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm2 = nn.RMSNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.SiLU(),
            nn.Linear(4 * d_model, d_model)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Self-Attention sub-layer
        norm_x = self.norm1(x)
        attn_out, _ = self.attn(norm_x, norm_x, norm_x)
        x = x + attn_out
        
        # Feed-Forward sub-layer
        norm_x2 = self.norm2(x)
        x = x + self.mlp(norm_x2)
        return x

class DeepTransformer(nn.Module):
    def __init__(self, num_layers: int, d_model: int, n_heads: int):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerBlock(d_model, n_heads) for _ in range(num_layers)
        ])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            # Execute layer forward pass with non-reentrant activation checkpointing
            x = checkpoint(layer, x, use_reentrant=False)
        return x

Under use_reentrant=False, PyTorch installs a pair of pack_hook and unpack_hook callbacks:

• pack_hook: Intercepts intermediate tensors scheduled for preservation by autograd and substitutes them with a placeholder or weak reference.
• unpack_hook: When autograd requests an activation during backward propagation, triggers re-execution of the specific sub-graph, restoring the required tensor directly into the active computation stream.

6. Compiler-Level Min-Cut Rematerialization (AOTAutograd)

With graph-capturing compilers like torch.compile and PyTorch 2.0 AOTAutograd (Ahead-Of-Time Autograd), activation checkpointing moves from manual module-level annotations to automated graph-partitioning algorithms.

AOTAutograd traces the joint forward-backward computational graph, constructing a directed acyclic graph (DAG) where nodes represent primitive operators and edges represent tensor values.

            [Input x]
            /       \
     [Operator A]    |
          |          |
     [Operator B]    |   <--- Min-Cut Partition Boundary
   ====================== (Tensors crossing cut are saved in memory)
          |          |
     [Operator C]    |
            \       /
            [Loss L]

The Min-Cut Objective

Let $G = (V, E)$ be the computational DAG, where $S \subset V$ represents forward operators and $T \subset V$ represents backward operators. The compiler formulates an optimization problem:

$$\min_{E_{cut}} \sum_{e \in E_{cut}} \text{Memory}(e) \quad \text{subject to} \quad \text{ComputeOverhead}(E_{cut}) \le \gamma \cdot C_{fwd}$$

Using max-flow min-cut formulations:

1. Every edge $e = (u, v)$ crossing from the forward sub-graph to the backward sub-graph represents an activation tensor saved in memory.
2. If an edge has a high memory cost but low recomputation cost (such as point-wise operations), the compiler cuts the edge and inserts a rematerialization node in the backward graph.
3. If an edge represents a compute-heavy GEMM result, the compiler preserves the tensor across the forward-backward boundary.

7. Interaction with FlashAttention, FSDP, and Model Parallelism

Activation checkpointing interacts directly with distributed parallelism strategies and fused kernel primitives:

FlashAttention Integration

Dao et al. (2022) in FlashAttention implemented online softmax tiling within GPU SRAM. Standard attention writes the $b \times a \times s \times s$ attention score matrix to HBM. FlashAttention never materializes this matrix in HBM during either forward or backward passes.

When FlashAttention is combined with selective activation checkpointing:

• FlashAttention eliminates the quadratic $\mathcal{O}(s^2)$ memory bottleneck at the kernel level via SRAM recomputation.
• Selective checkpointing eliminates linear projection and normalization memory overheads at the module level.
• Together, they allow scaling sequence lengths beyond $128\text{k}$ tokens on standard hardware clusters without incurring full $33\%$ recomputation penalties across all layers.

Fully Sharded Data Parallelism (FSDP / ZeRO-3)

When using FSDP, model parameters are un-sharded (via all-gather) immediately before forward execution and freed immediately after.

If full activation checkpointing is applied without coordination:

• Backward pass recomputation requires an additional all-gather collective communication phase to reconstruct sharded parameters a second time.
• Communication volume increases by $33.3\%$ per training step.

FSDP solves this via activation checkpointing wrappers (apply_activation_checkpointing) that hook directly into FSDP forward units, ensuring collective communication schedules overlap seamlessly with recomputation streams.

8. Summary Comparison of Activation Management Strategies

• Standard Forward-Backward:
• Peak Activation Memory: $\mathcal{O}(L \cdot b \cdot s \cdot h + L \cdot b \cdot a \cdot s^2)$
• Compute Overhead: $0.0\%$
• Best Used For: Small models, short sequence lengths ($s \le 2048$), and memory-unconstrained runs.
• Selective Activation Recomputation:
• Peak Activation Memory: $\mathcal{O}(L \cdot b \cdot s \cdot h)$ (eliminates non-GEMM and quadratic attention terms)
• Compute Overhead: $\approx 2\% - 4\%$
• Best Used For: Standard LLM pre-training and fine-tuning with FlashAttention where memory constraints are moderate.
• Uniform Full Checkpointing ($\sqrt{N}$ Segments):
• Peak Activation Memory: $\mathcal{O}(\sqrt{L} \cdot b \cdot s \cdot h)$
• Compute Overhead: $33.3\%$
• Best Used For: Deep transformer models and large batch sizes where static + activation memory exceeds available HBM.
• Checkpointing + Activation Offloading (CPU / NVMe):
• Peak Activation Memory: $\mathcal{O}(1 \text{ layer on GPU})$
• Compute Overhead: PCIe transfer latency overhead (mitigated by asynchronous pinned memory streams)
• Best Used For: Extreme long-context fine-tuning on consumer or single-node GPU hardware.

Sources

• Chen, T., Xu, B., Zhang, C., & Guestrin, C. (2016). Training Deep Nets with Sublinear Memory Cost. arXiv preprint arXiv:1604.06174.
• Korthikanti, V. A., Casper, J., Dey, S., Andersch, M., Shoeybi, M., & Catanzaro, B. (2022). Reducing Activation Recomputation in Large Transformer Models. arXiv preprint arXiv:2205.05198.
• Griewank, A., & Walther, A. (2000). Algorithm 799: Revolve: An Implementation of Checkpointing for the Reverse or Adjoint Mode of Computational Differentiation. ACM Transactions on Mathematical Software, 26(1), 19-45. DOI: 10.1145/347837.347846.
• Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems (NeurIPS 2022). arXiv:2205.14135.
• PyTorch Documentation. Activation Checkpointing and torch.utils.checkpoint. PyTorch Autograd Tutorials.