Does Activation Checkpointing change model convergence?

By design, it should not, but the forward replay must be equivalent to the original forward pass. Random numbers, global variables, device transfers, side effects, and non-deterministic kernels can all introduce differences, so you must compare loss, gradients, and parameter updates.

Why should I prefer use_reentrant=false?

PyTorch recommends the non-reentrant implementation. It supports early termination of unnecessary recomputation, keyword arguments, nested structures, and more complete Autograd scenarios, while also providing determinism_check and debug diagnostics.

Is Selective Recompute always faster than full block recomputation?

Not necessarily. It provides a larger strategy space, not an automatic guarantee. You must measure the memory-throughput Pareto frontier on your specific model, sequence length, parallelism, and operator backend to determine the optimal set of tensors to save and recompute.

LLM Activation Checkpointing in Production: Selective Recompute, Non-Reentrant Mode, and RNG Consistency to Reduce Training Memory

Background: Training Memory Isn’t Just Parameters

When deploying inference services, teams typically first focus on weights, KV Cache, and the memory allocator. But when you move to full fine-tuning or continued pre-training, it’s often activations that push memory to its limit first.

During the forward pass, Autograd needs to save a set of intermediate tensors for the backward pass to compute gradients. This memory footprint generally grows with micro-batch size, sequence length, hidden dimension, and number of network layers. You can think of the trend with this rough relationship:

activation_memory ∝ micro_batch × sequence_length × hidden_size × layers

This isn’t an exact formula—attention, MLP, intermediate projections, parallelism strategies, data types, and kernel implementations all change the constant factor—but it explains a common phenomenon: even after model parameters are sharded via FSDP or ZeRO, training can still OOM on activations as sequence length increases.

Activation Checkpointing works by saving fewer intermediate activations during the forward pass and re-executing the corresponding forward segment during the backward pass, trading memory pressure for extra computation. This isn’t about “how to persist training state” or checkpoint recovery after a failure; the “checkpoint” here refers to activation checkpoints in the computation graph.

Core Principle: Save Boundaries, Replay Internal Computation on Backward

Normal training saves all intermediate tensors needed for the backward pass. With Activation Checkpointing enabled, a wrapped region typically only retains its necessary inputs. When the backward pass needs internal activations, it re-executes the forward logic for that region.

This leads to three direct consequences:

Peak memory decreases: Many intermediate activations don’t need to live across the entire forward phase.
Computation increases: Discarded activations must be regenerated during the backward phase.
Correctness constraints tighten: The replay must be equivalent to the first forward pass in terms of control flow, random numbers, and device semantics.

Therefore, the engineering goal shouldn’t be “checkpoint all Transformer layers.” A more reasonable goal is:

Minimize recomputation overhead while meeting the target micro-batch size, sequence length, and peak memory budget, and prove that gradient semantics haven’t changed.

Don’t Treat Full Block Recomputation as the Only Option

Region-Level Activation Checkpointing

The most common approach is to divide the model into regions based on Transformer blocks, e.g., every 2, 4, or 8 layers forming one checkpoint region.

If a region is too large, the backward pass needs to recompute many expensive operators. If a region is too small, the overhead from boundary tensors, function calls, and scheduling increases. The optimal point depends on the model architecture, sequence length, FlashAttention backend, and micro-batch configuration—it can’t be determined by layer count alone.

Selective Activation Checkpointing

PyTorch’s Selective Activation Checkpointing (SAC) allows you to specify, within a single region:

Which operator outputs must be saved.
Which operators should be preferentially recomputed.
Which operators could be saved and then offloaded to CPU.

A practical starting point is: Preferentially save expensive matrix multiplications and attention results; recompute cheap element-wise operators, activation functions, and parts of normalization. This isn’t a fixed rule; you must ultimately verify with a profiler and memory snapshots.

The official PyTorch blog describes SAC as a fine-grained policy layer on top of standard Activation Checkpointing: standard AC recomputes all operations within a region, while SAC can avoid re-executing some high-cost operations.

Why Production Code Should Prefer Non-Reentrant Mode

PyTorch currently offers both reentrant and non-reentrant checkpoint implementations and explicitly recommends using use_reentrant=False.

Engineering advantages of non-reentrant mode include:

It can stop further useless recomputation early once the needed activations have been rebuilt.
Supports a more complete set of Autograd use cases.
Handles keyword arguments and more complex nested inputs.
Can enable determinism_check and debug to compare the shapes, types, devices, and operator traces of the original forward and the replay tensors.

Reentrant mode re-runs the entire function and has more restrictions on input/output requires_grad, nested structures, detach(), and backward APIs. When upgrading older projects, you can’t just change the default parameter; you must also verify custom Autograd Functions, gradient hooks, and module side effects.

RNG Consistency: Correct Dropout Doesn’t Guarantee Correct Random Numbers

Activation Checkpointing executes a forward segment an extra time. Without handling the random number generator (RNG) state, Dropout, stochastic depth, sampling routing, or other random operators might produce different results during replay, thus altering the gradients.

PyTorch saves and restores the RNG state by default, making the checkpointed version as deterministic as possible compared to the non-checkpointed version. This process has overhead, but you shouldn’t disable it prematurely when establishing a correctness baseline.

Pay special attention to three edge cases:

The default logic usually handles the CPU and one inferred device type; if a checkpoint region spans multiple device types, you must verify it separately.
If a function moves a tensor to a new device that wasn’t seen before, the framework can’t save that device’s RNG state in advance.
The RNG semantics of global random generators, custom CUDA kernels, or third-party operators might not be fully covered by the general mechanism.

Therefore, preserve_rng_state=False is only suitable when you’ve confirmed the region has no random operations, or when the business explicitly accepts changes in the random trajectory. It’s a performance option, not a default “optimization switch.”

Engineering Implementation: Converging from Full Block Recompute to Selective Strategies

Step 1: Establish a Baseline Without Recomputation

Record at least the following metrics, keeping the model, data batch, seed, parallelism, and mixed precision configuration fixed:

max_memory_allocated and max_memory_reserved.
Per-step latency and Tokens/s.
Time spent in forward, backward, and optimizer phases.
Loss, gradient norm, and count of non-finite values.
Peak memory for long and short sequence buckets.
Memory watermark and tail latency per rank.

Just recording “can it run?” isn’t enough. Activation Checkpointing often makes OOM disappear, but it can also stretch per-step time to an unacceptable level.

Step 2: Start with Full Transformer Blocks

Begin with an interpretable block granularity, e.g., enabling it only for a few layers in the middle or later part of the model, not the entire model at once. Gradually increase the number of covered layers to get a memory-throughput curve.

The NVIDIA Megatron Bridge performance guide distinguishes between recomputing only attention-related activations and recomputing the full Transformer block, allowing you to configure the number of recomputed layers. The official guide also notes that full block recomputation can significantly reduce activation memory but introduces substantial extra computation, so it should be considered a powerful but expensive option.

Step 3: Use Selective Recompute to Preserve Expensive Operators

Here’s a simplified PyTorch example. Operator names will change with PyTorch and Attention Backend versions; production code should include version compatibility tests.

from functools import partial
import torch
from torch.utils.checkpoint import (
    CheckpointPolicy,
    checkpoint,
    create_selective_checkpoint_contexts,
)

# Example: avoid recomputing high-cost matrix multiplications.
OPS_TO_SAVE = {
    torch.ops.aten.mm.default,
    torch.ops.aten.bmm.default,
    torch.ops.aten.addmm.default,
}

def policy_fn(ctx, op, *args, **kwargs):
    if op in OPS_TO_SAVE:
        return CheckpointPolicy.MUST_SAVE
    return CheckpointPolicy.PREFER_RECOMPUTE

sac_context = partial(create_selective_checkpoint_contexts, policy_fn)

def run_block(block, hidden_states, attention_mask):
    # Use a closure to fix keyword arguments, reducing model interface differences.
    def forward_fn(x):
        return block(
            hidden_states=x,
            attention_mask=attention_mask,
        )
    return checkpoint(
        forward_fn,
        hidden_states,
        use_reentrant=False,
        preserve_rng_state=True,
        determinism_check="default",
        context_fn=sac_context,
    )

Don’t just look at operator names when building your policy list. You need to verify on your target model:

Is Attention fused into a different ATen or custom operator?
Does torch.compile change graph partitioning and saving strategies?
Can the policy function recognize FlashAttention, Transformer Engine, or Triton kernels?
Does saving a particular output actually increase memory lifetime?

Step 4: Make “Correctness Replay” a Release Gate

Run three tracks with the same fixed batch of data:

Without Activation Checkpointing.
With standard Block Checkpointing.
With Selective Activation Checkpointing.

Compare the following:

Per-step Loss difference.
Maximum absolute error, relative error, and cosine similarity of gradients for key layers.
Weight difference after parameter updates.
Multiple fixed-seed replays with Dropout enabled.
Result consistency across single GPU, DDP, FSDP, or tensor parallelism.
Whether the error accumulates over multiple gradient accumulation micro-steps.

Numerical tolerance should be set based on BF16, FP16, FP8, operator backend, and parallel reduction order. You shouldn’t require bit-exact results across all configurations, but you also shouldn’t just compare the final loss.

Additional Risks in Distributed Training

All Ranks Must Use the Same Strategy

If some ranks enter a checkpoint region due to a conditional branch while others don’t, the order of collective communication operations during the backward pass can become inconsistent, potentially leading to an NCCL hang rather than a clear Python exception.

Avoid Side Effects Inside Replay Regions

The following logic should ideally be placed outside checkpoint regions:

Updating global counters.
Writing logs, caches, or queues.
Modifying external module state.
AllReduce or RPC calls that should only execute once.
Branches dependent on the current time, request ID, or non-deterministic external input.

The official PyTorch documentation explicitly warns: if the function call during backward replay differs from the original forward pass, it might cause an error or, worse, silently produce incorrect gradients.

Evaluate Pipeline Parallelism Per Stage

Different pipeline stages can have varying numbers of layers, activation volumes, and bubble ratios. A blanket “recompute N layers per stage” setting isn’t necessarily optimal. Measure peak memory and tail latency for each stage separately, then adjust the recomputation coverage.

How to Choose a Strategy

Training State	Preferred Solution	Reason
Need a small memory reduction to increase micro-batch	A few Block Checkpoints	Simple to configure and verify
Long sequences cause attention activations to dominate	Selective recomputation of attention	Targets the primary memory consumer
Full block recomputation causes too much throughput loss	SAC to preserve Matmul/Attention	Avoids re-executing expensive operators
Parameter memory is low after FSDP, but activations are still high	Compare AC vs. activation offload	Bottleneck has shifted from model state to activations
Many Dropout or random routing layers	Preserve RNG state and compare gradients	Prevents changes in random trajectory during replay
Model has complex nested inputs and custom gradients	Non-reentrant mode	Better compatibility and diagnostic capabilities

Common Misconceptions

Misconception 1: Memory reduction ratios can be directly copied from papers

Different models, sequence lengths, parallelism strategies, and kernel backends vary significantly. Classic research on selective recomputation reported substantial activation memory and execution efficiency gains for specific large-scale GPT training configurations, but these results can’t be directly expected for arbitrary fine-tuning tasks.

Misconception 2: Only look at GPU Utilization to judge gains

Recomputation might make the GPU busier, but Tokens/s can still drop. You need to observe per-step latency, effective token throughput, MFU, peak memory, and total training cost simultaneously.

Misconception 3: Disabling RNG saving is a free speedup

If the region contains Dropout or other random operations, disabling RNG state saving will change the replay path. You must first prove that the gradient difference is within an acceptable range.

Misconception 4: Activation Checkpointing and persistent Checkpointing are the same thing

The former is a computation-memory trade-off that happens between the forward and backward passes of each training step. The latter is for fault recovery and persisting training state, happening on the storage and restore path. They should be designed and monitored separately.

Misconception 5: Larger replay regions are easier

Larger regions generally save more memory, but they also increase computation and debugging costs. Production configurations should aim for the smallest recomputation set that meets the memory target.

Go-Live Checklist

References

PyTorch torch.utils.checkpoint documentation: https://docs.pytorch.org/docs/2.13/checkpoint.html
PyTorch Foundation, Current and New Activation Checkpointing Techniques in PyTorch: https://pytorch.org/blog/activation-checkpointing-techniques/
NVIDIA Megatron Bridge Performance Tuning Guide: https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-guide.html
Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198