What is the difference between activation checkpointing and gradient accumulation?

How to think about it

Activation checkpointing saves memory inside one forward/backward pass by discarding selected activations, the intermediate tensors needed to calculate parameter gradients, and recomputing them during backward. Gradient accumulation processes several micro-batches, small batches handled one at a time, before one optimizer update, the step that changes the weights, creating a larger effective batch, the total examples contributing to one update. The two techniques solve different problems and can be combined.

Why they solve different problems

A training step uses memory for model weights, their gradients, optimizer state such as Adam’s running averages, temporary workspaces, and activations. Activations are the values produced by layers during the forward pass. The backward pass uses many of those values to calculate how each weight contributed to the loss.

Without checkpointing, the framework keeps those values until backward finishes. A large transformer can therefore run out of memory even when the weights themselves fit comfortably.

Activation checkpointing changes that trade-off. It keeps only selected boundary values, discards internal values, and runs the relevant forward region again during backward. The recomputed values are then used to calculate gradients. You spend extra compute to save activation memory. The model’s weights, gradients, and optimizer state are still there.

Gradient accumulation changes the size and timing of the batch. Instead of processing 32 examples at once, you might process 4 examples, calculate their gradients, then process another 4, and continue until 32 examples have contributed. The framework adds those gradients into each parameter’s gradient buffer, but does not call optimizer.step() until the group is complete. The activations from each micro-batch can be released after its backward pass.

Common mistake: gradient accumulation does not retain the activations from all the micro-batches. It retains their accumulated gradients. If memory grows after every micro-step, something is probably retaining the computation graph, often retain_graph=True or a list of tensors that still have gradients attached.

A concrete 24 GB example

Suppose profiling a training job on a 24 GB GPU gives this deliberately concrete result:

• Model weights, gradient buffers, and optimizer state use 14.5 GB.
• CUDA workspaces and framework overhead use 2.0 GB.
• A micro-batch of 8 examples needs 7.0 GB for saved activations.

The same effective batch can be built in different ways:

Configuration	Effective batch	Saved activations at once	Profiled peak
8 examples, 4 micro-steps	32	7.0 GB	23.5 GB
4 examples, 8 micro-steps	32	3.8 GB	20.3 GB
4 examples, 8 micro-steps, checkpointed	32	2.0 GB	18.5 GB

The numbers are an example, not a universal benchmark. The point is the memory location being changed. Accumulation lets the job use a smaller micro-batch. Checkpointing reduces the saved activations inside that micro-batch. Neither technique removes the 14.5 GB occupied by weights, gradients, and optimizer state.

What the accumulation loop actually does

Assume loss_fn returns the mean loss for each equally sized micro-batch:

accumulation_steps = 8
optimizer.zero_grad(set_to_none=True)

for micro_step, batch in enumerate(loader):
    predictions = model(batch["inputs"])
    loss = loss_fn(predictions, batch["targets"])

    (loss / accumulation_steps).backward()

    if (micro_step + 1) % accumulation_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

Dividing the loss by accumulation_steps matters. Without it, eight mean losses are added together, producing gradients roughly eight times larger than the gradient of the corresponding 32-example mean loss. The first symptom may be unusually large gradient norms, unstable loss, or NaNs.

The code also clips after accumulation. Clipping each micro-batch separately is a different algorithm because it changes each gradient before the gradients are summed. In production, handle a final group containing fewer than eight micro-batches carefully: normalize by the actual group size, or by the total number of valid tokens when training a language model with padding and variable sequence lengths.

The nuance: a larger effective batch is not always identical

With a local micro-batch of size B, K accumulation steps, and W data-parallel workers, the usual effective batch is:

B_eff = B × K × W

That assumes each worker processes a distinct local micro-batch and the distributed gradients are averaged in the normal way.

For a simple model with an additive loss, correct normalization, no optimizer step between micro-steps, and no batch-dependent operations, accumulation can produce the same summed gradient as processing the whole effective batch at once. In real systems, the results are not perfectly identical.

Batch normalization updates statistics from each micro-batch, not from the entire effective batch. Dropout randomness can also differ. Gradient clipping, regularization, and mixed-precision scaling must happen at the intended point. A learning-rate scheduler should usually advance once per optimizer update, not once per micro-step. With eight micro-steps per update, advancing it on every micro-step changes the schedule relative to the number of examples seen.

For token-level losses, equal numbers of examples do not necessarily mean equal numbers of valid tokens. A micro-batch with 2,000 non-padding tokens should contribute more to a token-averaged loss than one with 1,000. Dividing every micro-batch loss by eight can therefore be wrong unless the batches have equal token counts.

Trade-offs and when not to use them

Use activation checkpointing when the memory bottleneck is saved activations and the job has enough compute budget for extra forward passes. It is especially useful for deep networks and long sequences. It generally has no value during ordinary inference because inference has no backward pass requiring those saved activations. It also cannot solve an out-of-memory error caused by model weights or optimizer state.

Use gradient accumulation when the optimizer should see a larger batch but the desired batch does not fit in memory. It is not free: several smaller forward/backward passes can have more kernel-launch overhead and worse GPU utilization than one well-sized batch. Very large effective batches also mean fewer optimizer updates for the same number of examples, which can change convergence and may require retuning the learning-rate schedule.

If the GPU has room for a full batch and the optimizer benefits from frequent updates, accumulation may make training slower without helping. If the model state itself consumes nearly all available memory, neither technique is the right first fix; parameter sharding, optimizer-state sharding, lower-precision state, or offloading addresses a different part of the memory budget.

What they’ll ask next

1. Is gradient accumulation exactly the same as increasing batch size?
   It can match the gradient of a larger batch when loss scaling and update timing are correct, but batch-normalization statistics, randomness, clipping, token weighting, and scheduling can make the training behavior different.

2. Can activation checkpointing and gradient accumulation be used together?
   Yes. Choose a micro-batch that gives acceptable GPU utilization, checkpoint the memory-heavy regions inside it, and accumulate enough micro-batches to reach the desired effective batch.

3. How do you decide which one to apply first?
   Profile the memory breakdown. If saved activations dominate, checkpointing attacks the cause. If a smaller micro-batch already fits but the optimizer needs more examples per update, accumulation is appropriate. If weights or optimizer state dominate, use a sharding or state-compression strategy instead.

Say this in the interview

“Checkpointing trades extra forward computation for less activation memory within one micro-batch, while gradient accumulation trades extra micro-steps and delayed updates for a larger effective batch; they are complementary, not interchangeable.”