Gradient Checkpointing Explained: Trade Compute for VRAM

Gradient checkpointing, explained in one sentence, trades a slice of GPU compute for a much bigger slice of GPU memory. Instead of storing every intermediate activation from the forward pass so the backward pass can use them to compute gradients, you keep only a fraction of them and recompute the rest on demand when the backward pass needs them. The training run gets slower, typically by 20-30%. The activation memory it needs can shrink from something that scales with the number of layers in your model to something that scales with roughly the square root of that number. That single trade is why a 70B fine-tune that would otherwise need two rented GPUs can run on one, and it's why turning gradient checkpointing on is really a decision about which GPU (and how many) you rent, not just a training flag you flip and forget.

TL;DR: Gradient Checkpointing Explained: Trade Compute for VRAM

Gradient checkpointing discards most forward-pass activations and recomputes them during the backward pass, trading extra compute for less VRAM.

• Memory: Activation memory can drop from scaling with layer count (n) to roughly sqrt(n), for one extra forward pass per mini-batch, per the 2016 sublinear memory paper.
• Compute cost: Full recomputation adds about 30% per-layer compute time, per NVIDIA's Megatron docs; selective recomputation cuts that overhead over 90% in one large-scale NVIDIA study.
• Enable it: torch.utils.checkpoint in PyTorch, or gradient_checkpointing=True in Hugging Face Trainer.
• Rental decision: Spheron H100 SXM5 runs $2.64/hr on-demand as of 17 Sep 2026; weigh that single-GPU rate against a two-GPU full-speed run. Compare live GPU pricing.

One thing to clear up before we go further: this is a different "checkpointing" from the kind that saves and restores training state after a spot instance gets reclaimed. If that's what brought you here, see spot GPU training resilience and preemption recovery instead. This post is about discarding and recomputing activations to save memory during a run that keeps going, not about surviving an interruption.

Why Activations, Not Weights, Are What Blow Up Training Memory

Training memory isn't dominated by the model's weights the way most people assume. Weights are fixed: a 70B-parameter model is roughly 140GB at FP16 whether you're running one training step or ten thousand. Gradients are the same size as the weights, and Adam's optimizer states typically double that again. Add those three up and you get a large but fixed number for a given model and precision. Our LLM VRAM requirements guide breaks down that weights-plus-gradients-plus-optimizer-states arithmetic in full if you need the exact figures for your model size.

Activations are the part that doesn't stay fixed. Every layer in the forward pass produces intermediate tensors, the outputs of attention, the outputs of each feed-forward projection, the normalized hidden states, and the backward pass needs every one of them to compute gradients through the chain rule. Unlike weights, activation memory scales with batch size × sequence length × hidden dimension × number of layers. Double your batch size or your context length and activation memory doubles right along with it, while your weight memory doesn't move at all. On a model with a long context window or a training batch pushed up for throughput, activations can end up being the single largest consumer of VRAM, not the parameters everyone sizes their GPU around first.

That's the problem gradient checkpointing exists to solve: it doesn't touch weights, gradients, or optimizer states at all. It only changes how much of that layer-by-layer, batch-and-sequence-scaling activation tensor you're forced to hold onto at once.

Gradient Checkpointing Explained: How It Trades Recompute for VRAM (With the Math)

The sqrt(n) Rule From the Original 2016 Paper

The technique traces back to a 2016 paper, "Training Deep Nets with Sublinear Memory Cost," which describes an algorithm that trains an n-layer network using memory that scales with the square root of n, rather than linearly with n, at the cost of one extra forward pass per mini-batch. The idea is straightforward once you see it laid out: instead of storing the activations for all n layers, you divide the network into roughly sqrt(n) segments and only checkpoint the activation at the boundary of each segment. During the backward pass, when you need the activations inside a segment, you re-run the forward pass for just that segment, starting from its checkpointed boundary tensor, to regenerate what you need.

The memory savings compound because you're never storing more than one segment's worth of intermediate activations plus the sqrt(n) checkpoint boundaries at any given time, instead of all n layers' worth. The price for that is a second forward pass through the checkpointed segments, which is where the extra compute cost comes from. That trade, a bit more wall-clock time for a much smaller memory footprint, is the entire idea, and it's why the technique shows up under both names "gradient checkpointing" and "activation recomputation" in different papers and frameworks.

What Actually Happens in the Backward Pass

In practice, most teams use PyTorch's built-in implementation rather than the original paper's exact algorithm, and it's worth understanding what it does mechanically. According to PyTorch's own documentation, torch.utils.checkpoint does not store the intermediate tensors produced inside the checkpointed function. Instead, it re-invokes that function during the backward pass to regenerate the tensors it needs for gradient computation, discarding them again immediately after.

There are two implementations under the hood, reentrant and non-reentrant, selected by the use_reentrant argument. PyTorch recommends use_reentrant=False for new code. The reentrant version always fully recomputes the checkpointed segment's forward pass before backward can proceed. The non-reentrant version records the autograd graph as it goes and can stop recomputation as soon as every intermediate activation actually needed for the current gradient computation has been regenerated, which is generally more memory-efficient and more compatible with other autograd features like torch.autograd.grad. In code, this looks like:

import torch
from torch.utils.checkpoint import checkpoint

class TransformerBlock(torch.nn.Module):
    def forward(self, x):
        return checkpoint(self._forward_impl, x, use_reentrant=False)

    def _forward_impl(self, x):
        x = self.attention(x)
        x = self.mlp(x)
        return x

Nothing about the model's math changes. The gradients computed with checkpointing are numerically identical to gradients computed without it (modulo floating-point nondeterminism from re-running kernels). What changes is when the GPU is asked to hold which tensor in memory.

Full vs Selective Recomputation: Not All Activations Cost the Same to Save

Checkpointing every layer is the blunt version of this technique, and it isn't the version most production training setups actually want.

The reasoning behind selective recomputation is that not every activation tensor is equally expensive to keep around, and not every activation tensor is equally cheap to recompute. Some tensors, like the attention softmax output, are large in memory but cheap to regenerate. Others are small but sit on the critical path of an expensive matmul. Selective recomputation picks out the specific tensors that free the most memory per unit of recompute cost, and leaves the rest stored, rather than treating every activation as equally disposable.

The clearest evidence for how much this matters comes from "Reducing Activation Recomputation in Large Transformer Models", a 2022 NVIDIA paper that combined selective activation recomputation with sequence parallelism. Tested on a 530B-parameter GPT-3-style model trained across 2,240 A100 GPUs, the paper reports that this combination reduces activation memory by 5x, while reducing the execution time overhead from recomputation by over 90% compared to full recomputation. In throughput terms, the same paper reports that selective recomputation improves training efficiency from 42.1% to 54.2% of theoretical peak FLOPs on that 530B model, roughly a 29% speed difference for the same hardware footprint. That gap, 42.1% versus 54.2% model FLOPs utilization, is the practical cost of choosing full recomputation when selective recomputation was available and would have hit the same memory target.

The takeaway for anyone configuring this today: if your framework offers a selective or partial recomputation mode, reach for it before reaching for full checkpointing on every layer. Full recomputation is the easy, all-or-nothing lever. Selective recomputation is the one that actually shows up as the difference between 42% and 54% GPU utilization on the same job.

Enabling It: PyTorch, Hugging Face Trainer, and Megatron/DeepSpeed

The mechanics differ by framework, but the concept is identical everywhere: turn activation recomputation on for the layers where it's worth the trade, and leave it off where you have memory headroom.

Raw PyTorch. As shown above, wrap the module you want checkpointed with torch.utils.checkpoint.checkpoint(fn, *args, use_reentrant=False). This is the lowest-level control and the one every higher-level framework builds on.

Hugging Face Trainer. The simplest path for fine-tuning jobs:

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="./out",
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
)

Or call it directly on a loaded model with model.gradient_checkpointing_enable(). Hugging Face's Transformers documentation also describes an every_n_layers argument for partial checkpointing, which lets you checkpoint, say, every second or third transformer block instead of an all-or-nothing choice across the whole model, giving you a dial between full memory savings and full speed rather than a single switch.

Megatron and DeepSpeed. At cluster scale, both frameworks expose recomputation as a first-class config option rather than a model-code change. Megatron's activation-recomputation settings let you choose full or selective recomputation per transformer layer, matching the distinction covered above. DeepSpeed's activation checkpointing config (under the activation_checkpointing block in a ZeRO config) offers similar granularity, plus CPU offload of checkpointed activations for teams that want to trade host RAM for even more GPU memory headroom.

Fine-tuning frameworks. If you're using a higher-level trainer rather than writing your own loop, gradient checkpointing is usually a one-line YAML or config flag. Our Axolotl vs Unsloth vs TorchTune comparison covers how each framework exposes this setting, and for mixture-of-experts models specifically, our guide to fine-tuning MoE LLMs covers how expert-routing tensors change the activation-memory profile checkpointing has to work against.

One interaction worth knowing about: kernel-fusion libraries like Liger Kernel reduce activation memory by eliminating intermediate tensors at the operator level rather than by recomputing them, which is a complementary technique, not a competing one. Liger's own documentation flags a specific gotcha here: gradient checkpointing and certain fused kernels need to be applied in the right order relative to each other, so check that guide before stacking the two.

The GPU Rental Decision: Fewer GPUs at ~20-30% Slower vs More GPUs at Full Speed

This is where the math above turns into an actual invoice. Say you're fine-tuning a model that just barely doesn't fit in one GPU's VRAM at full speed, no checkpointing, but does fit with room to spare once checkpointing is turned on. You now have two real options: rent two GPUs and run at full speed with tensor or pipeline parallelism splitting the model across them, or rent one GPU, turn on gradient checkpointing, and accept a run that takes roughly 20-30% longer.

The naive comparison, "two GPUs run 2x as fast, so it's a wash," is wrong in both directions. Multi-GPU training doesn't scale linearly because of communication overhead between devices, a point our GPU count versus training speedup piece covers in detail: doubling GPU count typically buys you noticeably less than double the throughput, because gradient synchronization and cross-device communication time doesn't shrink the way compute time does. So the real comparison usually looks more like this:

Run the actual numbers rather than trusting either row on instinct. On Spheron, H100 SXM5 is priced at $2.64/hr on-demand and $2.09/hr on spot as of 17 Sep 2026, billed per minute with no minimum rental period, so a run that finishes 25% later than a theoretical full-speed baseline doesn't get penalized by hourly rounding the way it might on a provider that bills in whole-hour blocks. That per-minute billing matters more here than it sounds: the entire value of the single-GPU-plus-checkpointing path is that it lets a slower, cheaper run finish and get charged for exactly the extra minutes it used, not a rounded-up extra hour. Because bare-metal instances aren't shared with other tenants, the 20-30% recompute overhead you measure is attributable to your own workload, not muddied by a noisy neighbor. Spot capacity is worth pricing into the same comparison: a checkpointed run that's already tolerant of running slower is usually also tolerant of the occasional preemption and resume that spot carries. Spot and on-demand track separate pools of supply, though, so don't assume one is automatically cheaper than the other; check both rates for your GPU on the day you're actually renting and run the same GPU-hours math against whichever number you'd pay.

Pricing fluctuates based on GPU availability. Spheron rates above are live as of 17 Sep 2026; other providers reflect their most recent published rates and may have changed. Check current GPU pricing → for live rates.

Our 12-person startup case study is a concrete example of the same underlying discipline: sizing a training run against a measured cost curve instead of assuming more hardware is automatically the faster or cheaper choice. The general rule from that comparison table holds here too: one GPU running longer is often cheaper than two GPUs running shorter, once you price both paths against the same live rates rather than assuming linear scaling in either direction.

There's a real limit at the other end of this trade, though. Once your model genuinely needs multi-node training regardless of checkpointing, tensor and pipeline parallelism across nodes, the interconnect fabric between nodes starts to matter as much as GPU count. Confirm what fabric a multi-node cluster actually gives you before assuming near-linear scaling; Spheron's documentation covers what's available by default versus what needs to be requested for a custom cluster.

When to Skip Gradient Checkpointing

Gradient checkpointing is not a free win you should always leave on, and treating it as a default setting rather than a deliberate trade-off wastes GPU-hours in the other direction. Skip it, or turn it off, when:

• You already have VRAM headroom. If your batch size and sequence length comfortably fit without checkpointing, turning it on just adds recompute cost for no benefit. Check your actual memory usage with nvidia-smi under load before assuming you need it.
• The run is short. A quick LoRA fine-tune that finishes in minutes rarely justifies the added complexity of tuning checkpointing granularity; just rent the GPU tier that fits without it.
• Your workload is already compute-bound, not memory-bound. If your GPU is already near 100% utilization without checkpointing, adding an extra forward pass on top of that only slows you down; it doesn't unlock anything, because memory wasn't your constraint.
• You're doing inference, not training. Gradient checkpointing exists to reduce memory needed for the backward pass and gradient computation. Inference has no backward pass, so there's nothing for it to save; the memory levers that matter for serving are different ones like KV cache management and PagedAttention, covered in our LLM VRAM calculator for Llama 4.
• You're already multi-node with room to spare on interconnect. If you've committed to a cluster large enough that per-GPU memory was never the constraint, spending 20-30% more compute to save memory you don't need is a straightforward loss.

The decision, in short, comes down to whether the GPU-hours you'd spend on recompute cost less than the GPU-hours (and GPU count) you'd save by fitting a bigger effective batch or a bigger model onto less hardware. That's a number worth actually computing against current prices before you commit a training budget to either path.

Whether the answer is one GPU with checkpointing turned on or several running at full speed, the fastest way to find out is to price both paths against real rates instead of assumptions. Spheron's bare-metal H100, H200, A100, and B200 instances bill per minute with no minimum commitment, so testing a checkpointed run's actual wall-clock cost doesn't lock you into anything longer than the run itself.

Rent H100 GPU →