Gradient checkpointing
Gradient checkpointing trades extra compute for lower GPU memory by recomputing activations in the backward pass. Learn how it works, costs and best practices.
Overview
Gradient checkpointing is a training technique that lowers GPU memory use by storing only a few intermediate results (called activations) during the forward pass and recomputing the rest when the backward pass needs them. It is also called activation checkpointing, activation recomputation or rematerialization. The price is extra compute: a typical setup adds roughly one extra forward pass per training step, which often means a slowdown in the range of twenty to thirty percent. The reward is large. Memory for activations can drop from growing linearly with network depth to growing roughly with the square root of depth, which lets you train bigger models, use longer sequences or use larger batches on the same hardware.
If your training job fails with an out of memory error and the model weights themselves fit comfortably, activations are usually the culprit, and gradient checkpointing is usually the first remedy to try. It is a standard feature in major deep learning libraries, it needs only a few lines of code, and it does not change the math of training. The gradients you get are the same as without checkpointing, up to tiny numerical differences. For models that read whole tables as sequences of cells, such as the tabular foundation models described on /model/architecture, it is one of the main tools that make long table contexts trainable.
What it is
Gradient checkpointing is a memory-saving strategy for backpropagation. To understand it, you first need to know why training uses so much memory.
The memory bill of training
When you train a neural network, memory goes to four main things:
- Parameters: the weights of the model.
- Gradients: one value per parameter, produced by the backward pass.
- Optimizer state: for example, two extra values per parameter for Adam style optimizers.
- Activations: the outputs of every layer, saved during the forward pass so that the backward pass can use them.
The first three scale with model size. The fourth scales with model depth, batch size and sequence length. For transformers, activations can easily become the largest item, especially when the input is long. Attention layers also produce score matrices whose size grows with the square of the sequence length, unless you use a fused kernel such as the one described in FlashAttention.
The idea in one sentence
Instead of keeping every activation, keep only some of them (the checkpoints), and when the backward pass needs a missing activation, recompute it from the nearest earlier checkpoint.
A plain analogy
Imagine you are walking a long trail and want to be able to retrace your steps exactly. One option is to drop a breadcrumb at every step. That takes a huge bag of breadcrumbs, but retracing is instant. Another option is to drop a breadcrumb only every hundred steps. Your bag stays light. When you walk back, you reach a gap and have to re-walk the last stretch from the previous breadcrumb to remember where each step was. You spend more time, but you carry far less. Gradient checkpointing is the second strategy.
What it is not
Gradient checkpointing does not shrink the model, does not change the loss and does not change the optimizer. It is not quantization and it is not pruning. It only changes which intermediate values are kept in memory and which are rebuilt on demand.
Where it came from
The idea of trading computation for memory in reverse mode automatic differentiation is old. Researchers in scientific computing studied checkpointing schedules for adjoint computations long before deep learning became popular. Those methods asked a clean question: given a chain of operations and a limited number of saved states, which states should you save so that the total recomputation is minimal?
For deep learning, a widely cited paper from 2016, titled "Training Deep Nets with Sublinear Memory Cost", brought the idea to neural networks. It showed that a network with n layers could be trained with memory cost on the order of the square root of n, at the price of one extra forward computation. That result made very deep networks practical on limited GPUs.
Since then the technique has become routine. Deep learning frameworks ship built in helpers, and large model training stacks, including DeepSpeed ZeRO and FSDP, treat activation checkpointing as a standard option. As transformer models grew and context windows lengthened, the technique moved from a niche trick to a default setting for many large training runs.
Related lines of work
- Selective recomputation: instead of recomputing whole blocks, recompute only cheap but memory heavy operations, such as activation functions, dropout masks and attention softmax parts, and keep the expensive matrix multiplications.
- Offloading: move saved activations to CPU memory and bring them back when needed. This trades memory for transfer time instead of compute.
- Reversible layers: design layers whose inputs can be recovered exactly from their outputs, so nothing needs to be stored.
- Fused attention kernels: avoid storing the large attention matrix by recomputing pieces of it inside the backward kernel. This is a form of checkpointing built into the kernel itself.
How it works
Step 1: the normal forward and backward pass
In a standard training step, the forward pass runs layer 1, then layer 2, and so on up to the loss. Each layer saves what it will need later. The backward pass then starts at the loss and walks backward, using the saved values to compute gradients for every layer.
Memory used by activations is proportional to the number of layers, because every layer holds something.
Step 2: choose checkpoints
With gradient checkpointing, you divide the network into segments. A segment might be one transformer block, or a group of several blocks. You decide that only the input to each segment will be saved. Everything computed inside a segment is thrown away after the segment finishes its forward pass.
Step 3: forward pass with fewer saves
The forward pass runs as before, but it runs the segments in a mode where no inner activations are recorded for the backward pass. After each segment, only the segment output (which is the input to the next segment) is kept.
Step 4: backward pass with recomputation
When the backward pass reaches a segment, it first reruns the forward computation for that segment from its saved input. This time it keeps the inner activations, because it needs them right now. It then runs the backward computation for that segment, produces gradients, and immediately frees the inner activations. Then it moves to the previous segment.
So at any moment, memory holds the saved segment inputs plus the activations for one segment being processed. That is far smaller than all activations at once.
Step 5: the memory and compute math
Suppose a network has n layers and you split it into k segments of n over k layers each.
- Saved segment inputs: k values.
- Inner activations alive at once: n over k.
- Total activation memory: k plus n over k.
This is smallest when k is about the square root of n, giving a memory cost near two times the square root of n, compared with n without checkpointing.
For compute, each layer is run forward once in the original pass and once more during recomputation. The backward pass costs about twice a forward pass. So a normal step costs about three forward units (one forward plus two for backward), and with checkpointing it costs about four. That is the origin of the common estimate of one third extra compute, though in practice real slowdowns are often lower, because not everything in a step is layer compute.
An everyday analogy for the schedule
Think of a cook following a long recipe with many stages, who must later explain each stage to a friend in reverse order. Writing down every detail of every stage takes lots of paper. Instead, the cook writes down only the state of the dish at the end of each major stage. When asked to explain stage four, the cook re-cooks stage four from the written state of stage three, explains it, and throws away the extra notes. Paper use stays low. The cook works a bit more.
Randomness and determinism
Recomputation must reproduce the same values as the original forward pass. Layers with randomness, such as dropout, need special care. Good implementations save and restore the random number generator state, so the recomputed dropout mask matches the original. If you write custom code, check that anything stochastic is handled in this way.
Interaction with other pieces
- Batch normalization: layers that update running statistics in the forward pass may update them twice if recomputed. Transformers using layer normalization do not have this issue.
- Mixed precision: works fine. Recomputation uses the same precision settings as the original pass. See mixed precision.
- Distributed training: checkpointing is local to each device, so it combines freely with data parallelism, tensor parallelism and pipeline parallelism.
Key concepts and terms
| Term | Meaning |
|---|---|
| Activation | The output of a layer for a given input, kept so that gradients can be computed later. |
| Forward pass | Running the input through the network to get predictions and loss. |
| Backward pass | Computing gradients of the loss with respect to every parameter, from the end to the start. |
| Checkpoint | A saved activation that serves as a restart point for recomputation. Not the same as a saved model file. |
| Segment | A group of consecutive layers treated as one unit for checkpointing. |
| Recomputation | Running part of the forward pass a second time during the backward pass. |
| Rematerialization | Another word for recomputation, common in compiler literature. |
| Peak memory | The maximum memory used at any moment during a training step. This is what causes out of memory errors. |
| Sublinear memory | Memory that grows more slowly than the number of layers. |
| Selective checkpointing | Recomputing only certain cheap operations rather than whole blocks. |
| Offloading | Moving saved tensors to slower memory such as CPU RAM instead of recomputing them. |
| Throughput | Samples or tokens processed per second. Checkpointing lowers it slightly. |
A note on naming: "checkpoint" also means a saved copy of model weights on disk. Gradient checkpointing has nothing to do with that meaning. Context usually makes this clear, but it is a common source of confusion for newcomers.
A worked example
A numeric walk-through
Take a toy network with 16 layers. Each layer produces an activation that takes 1 unit of memory.
Without checkpointing:
- Activations stored: 16 units.
- Compute per step: 1 forward plus 2 backward, about 3 units of work per layer.
With checkpointing in 4 segments of 4 layers:
- Segment inputs saved: 4 units.
- During backward, one segment at a time is rebuilt: up to 4 units alive.
- Peak activation memory: about 8 units, half of the original.
- Compute per step: 1 forward plus 1 recompute plus 2 backward, about 4 units per layer.
So memory fell by half, and compute rose by a third. For a 100 layer network with 10 segments, memory falls from 100 to about 20 units, a fivefold saving, for the same one third extra compute. The deeper the network, the better the trade.
Code in PyTorch
The simplest way to use the technique in PyTorch is the checkpoint helper, which wraps a function or module call.
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
class Block(nn.Module):
def __init__(self, dim):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(dim, 8, batch_first=True)
self.ff = nn.Sequential(
nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)
)
def forward(self, x):
h = self.norm(x)
a, _ = self.attn(h, h, h, need_weights=False)
x = x + a
return x + self.ff(self.norm(x))
class Net(nn.Module):
def __init__(self, dim=512, depth=24, use_ckpt=True):
super().__init__()
self.blocks = nn.ModuleList([Block(dim) for _ in range(depth)])
self.use_ckpt = use_ckpt
def forward(self, x):
for blk in self.blocks:
if self.use_ckpt and self.training:
x = checkpoint(blk, x, use_reentrant=False)
else:
x = blk(x)
return x
model = Net().cuda()
x = torch.randn(8, 2048, 512, device="cuda", requires_grad=False)
loss = model(x).pow(2).mean()
loss.backward()
print(torch.cuda.max_memory_allocated() / 1e9, "GB peak")
Run this twice, once with use_ckpt=True and once with use_ckpt=False, and compare the printed peak memory and the step time. You will normally see a clear drop in peak memory and a modest rise in time.
Library level switches
Many high level libraries expose a single flag. In the Hugging Face ecosystem, for example, a model method named gradient_checkpointing_enable turns the technique on for supported transformer models, and a training argument can do the same. In DeepSpeed ZeRO and FSDP, there are configuration options to wrap blocks with activation checkpointing. Check the documentation of the version you use for the exact names.
How to measure the effect
- Fix the batch size and sequence length.
- Reset the peak memory counter before the step.
- Run one training step and record peak memory and wall time.
- Repeat with checkpointing on.
- Compare. Then increase batch size or sequence length with checkpointing on until you hit the limit again. That shows the real benefit.
The real benefit is not just lower memory at the same settings. It is the larger batch or longer sequence that now fits.
Strengths and limitations
Strengths
- Large memory savings: activation memory can fall by a large factor, especially for deep models.
- Mathematically equivalent: gradients match those of normal training, apart from tiny floating point differences.
- Easy to adopt: often one line of code or one configuration flag.
- Composable: works with mixed precision, distributed strategies and most optimizers.
- Enables longer context: sequence length is a major driver of activation memory, so checkpointing directly extends the context you can train on.
- Enables bigger batches: larger batches can improve hardware utilization, which can offset part of the compute cost.
Limitations
- Extra compute: expect a slowdown, commonly in the range of ten to thirty percent or more depending on the model and the setup.
- Does not reduce parameter, gradient or optimizer memory: if those do not fit, you need sharding or lower precision instead.
- Granularity matters: wrapping too small or too large a unit gives poor trade-offs.
- Randomness and side effects: layers with stateful behavior need careful handling.
- Debugging is harder: stack traces and profilers show repeated work, and some tools get confused by recomputation.
- Not free on already compute bound jobs: if you are limited by compute rather than memory, the slowdown is a real cost with no gain, so only switch it on when memory is the bottleneck.
When it helps most
| Situation | Benefit |
|---|---|
| Deep transformer, long sequences | High |
| Large batch needed for stable optimization | High |
| Small model on a large GPU | Low, not needed |
| Memory limit caused by optimizer state | None, use sharding |
| Inference only | None, no backward pass |
How it compares with alternatives
Gradient checkpointing is one of several ways to reduce training memory. They attack different parts of the bill and are often combined.
| Method | What it saves | Cost | Typical use |
|---|---|---|---|
| Gradient checkpointing | Activations | Extra forward compute | Long sequences, deep models |
| Mixed precision | Activations and weights copies | Small risk of numeric issues | Almost always on |
| FlashAttention | Attention score matrices | None, usually faster | Any transformer |
| Gradient accumulation | Activations (smaller micro batches) | More steps per update | Emulate large batches |
| DeepSpeed ZeRO / FSDP | Parameters, gradients, optimizer state | Communication | Large models across many GPUs |
| Activation offloading | Activations | Transfer time | When recompute is too slow |
| bitsandbytes 8 bit optimizers | Optimizer state | Small accuracy risk | Single GPU fine tuning |
| LoRA / QLoRA | Gradients and optimizer state | Limited capacity | Fine tuning |
Checkpointing versus gradient accumulation
Both let you work with big effective batches. Gradient accumulation splits a batch into micro batches and sums the gradients, so activations of only one micro batch exist at a time. It costs no recompute but needs more passes and can hurt hardware utilization when micro batches get tiny. Checkpointing keeps the batch whole but throws away inner activations. Many teams use both.
Checkpointing versus offloading
Offloading moves activations to CPU memory. It avoids recomputation but depends on transfer bandwidth. On machines with fast links it can be attractive. On others, recomputation is cheaper than waiting for transfers.
Checkpointing versus model parallelism
Splitting a model over several devices, as in pipeline parallelism or tensor parallelism, spreads memory across hardware but needs more hardware. Checkpointing makes better use of what you already have. In large runs the two are used together.
Role in tabular foundation models
Tabular foundation models are pre-trained on many tables so that they can make predictions on a new table without task specific training. Their inputs look different from text. A table has rows and columns, mixed types (numbers, categories, dates, short text), missing values, and often columns with no meaningful names. Gradient checkpointing matters for these models in several concrete ways.
Tables become long sequences
Many tabular foundation models, such as TabPFN and TabICL, treat a whole dataset as the context. Labeled rows and the query rows are processed together, so the model can perform in-context learning: it reads the training rows and predicts the query rows in a single forward pass, without gradient updates at prediction time.
That makes the effective sequence long. A table with a few thousand rows and a few dozen columns contains tens of thousands of cells. If a model attends over rows, over columns, or over both (an approach often called two dimensional or axial attention, related to axial attention), activation memory grows quickly with the number of rows and columns. During pre-training, you want to show the model large tables so it learns to use many examples as context. Checkpointing is what lets those large tables fit.
Row and column structure
Models that respect table structure often apply attention across columns inside a row, and then across rows inside a column. Each stage creates its own activations. Stacking many such layers multiplies memory use. Wrapping each row attention block and each column attention block in a checkpoint is a natural way to cap peak memory. The segment boundaries line up with the architecture, which makes the trade-off efficient.
Mixed types and missing values
Cell encoders for numbers, categories and missing flags add extra layers at the input. They are cheap in parameters but produce one activation per cell, and the number of cells is large. Because these early activations are numerous, they are good candidates for recomputation. Recomputing a small encoder is cheap, and the saved memory can be large.
No column names
When column names are absent or unreliable, a model has to infer the role of a column from its values alone. This tends to push designs toward deeper networks with more cell level processing, since the model cannot lean on text names. Deeper models increase the benefit of checkpointing.
Pre-training on synthetic and real tables
Pre-training often uses very large numbers of synthetic tables drawn from random generators (see synthetic data generation), plus real tables. Training runs are long, so a thirty percent slowdown matters, but the larger tables or batches that checkpointing allows can improve learning per step. Teams usually benchmark both settings and pick the cheaper path to a given quality. A curriculum that moves from small simple tables to larger complex ones can also reserve checkpointing for the later, larger stages only, avoiding overhead early on.
Continual learning, imputation and anomaly detection
- Continual learning: when a pre-trained model is updated on new table families, memory headroom lets you train on larger tables without changing the recipe.
- Imputation: masked cell objectives, similar to masked cell modeling, predict hidden cells from the rest of the table. They process full tables with many masked positions, so memory is dominated by activations over all cells.
- Anomaly detection: scoring rows by how surprising they are under a learned model benefits from long contexts of normal rows. See anomaly detection.
- Sector and semantic classification: heads that classify a table into a domain or sector read pooled column and row representations, and the backbone below them is where memory is spent.
For a broader picture of the architecture family, see /model/architecture, and for how such models are evaluated, see /model/benchmarks and /model/research.
A practical note
Inference of an in-context tabular model needs no backward pass, so checkpointing does not apply there. It is purely a training and fine tuning tool. If you only use a pre-trained model for prediction, you will not need it. If you pre-train, fine tune or run gradient based adaptation on large tables, you probably will.
Estimating activation memory for a table model
A rough estimate helps you decide where to place checkpoints. For a block that attends over cells, the saved tensors per layer are roughly proportional to:
- the batch size,
- the number of cells in the table context (rows times columns),
- the hidden width of the model,
- a constant that depends on how many intermediate tensors the block keeps (for a standard transformer block, a dozen or so).
Multiply by the number of layers and by the bytes per value, and you get the activation bill. Because the number of cells is rows times columns, doubling both rows and columns quadruples the cell count. That is why a table of modest looking shape can overwhelm memory.
| Change to the input | Effect on activation memory | Effect of checkpointing |
|---|---|---|
| Double the rows | About double | Keeps growth but at a much lower level |
| Double the columns | About double | Same |
| Double rows and columns | About four times | Same, lower level |
| Double the depth | About double | Grows with the square root instead |
| Double the batch size | About double | Same, lower level |
Attention layers that compare every cell with every other cell add a term that grows with the square of the cell count. Row wise and column wise attention reduce that term, since each cell only attends within its own row or column. Fused kernels remove the stored score matrix. Checkpointing then handles whatever remains.
Selective recomputation in more detail
Not all operations are equal. A large matrix multiplication is expensive to repeat but its output is not especially large. An activation function, a dropout mask or a layer normalization is cheap to repeat but its output is as large as its input. Selective recomputation exploits this difference:
- Keep the outputs of the expensive operations.
- Throw away the outputs of the cheap, memory heavy operations.
- Recompute only the cheap ones during the backward pass.
The result is a large share of the memory saving for a small share of the extra compute. Several large training stacks offer this mode, and it is worth trying when full block recomputation costs too much throughput.
Design tips for new table architectures
- Group layers into blocks with a clear single input and single output, so that checkpoint boundaries are simple.
- Avoid hidden state that changes in the forward pass, such as running counters, inside a checkpointed block.
- Keep cell encoders and decoders in their own modules, so they can be checkpointed or left alone separately.
- Make dropout and random masking use the framework random state, so recomputation reproduces them.
- Test with a tiny model first. Compare gradients with and without checkpointing on a few batches. They should agree closely.
Checking gradient agreement
A quick test gives confidence before a long run:
import copy, torch
def grads(model, x, ckpt):
m = copy.deepcopy(model)
m.use_ckpt = ckpt
m.zero_grad()
m(x).pow(2).mean().backward()
return [p.grad.clone() for p in m.parameters()]
torch.manual_seed(0)
model = Net(dim=64, depth=4).cuda().train()
x = torch.randn(2, 32, 64, device="cuda")
g_a = grads(model, x, False)
g_b = grads(model, x, True)
print(max((a - b).abs().max().item() for a, b in zip(g_a, g_b)))
The printed difference should be tiny, at the level of floating point noise. If it is large, something in a block is not reproducible on recomputation.
Practical guidance and best practices
Decide whether you need it
- Check where memory goes. Use the framework memory summary or a profiler.
- If activations dominate peak memory, turn checkpointing on.
- If parameters and optimizer state dominate, look at sharding, lower precision optimizers or parameter efficient fine tuning first.
Choose the right granularity
- Start by checkpointing each transformer block. This is a good default.
- If memory is still too high, group two blocks per checkpoint or checkpoint sub-blocks.
- If compute overhead is too high, checkpoint only every second or third block.
- Consider selective recomputation, where available, to skip the expensive matrix multiplications.
Combine with other techniques
- Use mixed precision first. It is nearly free and lowers activation memory.
- Use a fused attention kernel such as FlashAttention so you do not store attention matrices.
- Add checkpointing next.
- Then consider sharding with FSDP or DeepSpeed ZeRO when parameters and optimizer state no longer fit.
Spend the savings well
After turning checkpointing on, raise the batch size or sequence length until memory is nearly full again. If you do not use the freed memory, you only paid the compute cost and got nothing back.
Check correctness
- Compare the loss curve for the first few hundred steps with and without checkpointing. They should match closely.
- Verify that random seeds and dropout behave as expected.
- Use the non reentrant implementation if your framework offers both, since it handles more cases and gives clearer errors.
Measure throughput, not just memory
Record tokens or samples per second. A setting that saves memory but halves throughput may not be what you want. Often a moderate setting gives most of the memory gain for a small fraction of the cost.
Keep inputs requiring gradients in mind
In some implementations of the older reentrant style, if none of the inputs to a checkpointed segment require gradients, the segment may not receive gradients at all. This is a well known pitfall at the first layer when the input is raw data. The newer non reentrant style avoids this.
Quick configuration checklist
training:
mixed_precision: bf16
fused_attention: true
gradient_checkpointing:
enabled: true
granularity: block # block | every_n_blocks | selective
every_n: 1
micro_batch_size: 8
gradient_accumulation_steps: 4
sharding: fsdp # only if parameters or optimizer state do not fit
The exact keys will differ by framework, but the decisions are the same.
Common mistakes
- Turning it on everywhere by default. If memory is not a problem, you pay extra compute for nothing.
- Forgetting to use the freed memory. The benefit comes from a larger batch or longer context.
- Checkpointing the wrong unit. Wrapping a single tiny layer saves little and adds overhead. Wrapping the whole model saves little either, since the peak is then the full recomputation.
- Ignoring randomness. Custom layers with dropout or random augmentation that do not restore random state can give inconsistent gradients.
- Using it at inference. There is no backward pass at inference, so it does nothing useful.
- Expecting it to fix parameter memory. It does not touch weights, gradients or optimizer state.
- Confusing it with model checkpoints. Saving weights to disk is a different concept.
- Mixing it with layers that mutate state. Batch statistics or counters updated in the forward pass can be updated twice.
- Not checking the output with
use_reentrantchoices. Some versions warn about the default. Pick the non reentrant style explicitly. - Comparing runs without fixed seeds. Small differences in loss curves can be noise, not a bug. Fix seeds and compare early steps.
- Neglecting data loading. If the input pipeline is the bottleneck, extra compute from recomputation may be hidden, which can make the cost look smaller than it is on faster hardware.
Frequently asked questions
What is gradient checkpointing in simple terms?
It is a way to save GPU memory during training by not keeping every intermediate result from the forward pass. Instead, the training code keeps a few key results and recomputes the others when it needs them for the backward pass. You pay with some extra computing time and gain a lot of memory.
Does gradient checkpointing change the trained model?
No, not in any meaningful way. It computes the same gradients as normal training, because the recomputed activations are the same values as the ones that were thrown away. Tiny differences can appear because of floating point behavior and non deterministic kernels, but they are of the same size as normal run to run noise.
How much slower is training with gradient checkpointing?
The theoretical cost is one extra forward pass per step, which is about one third more compute for a standard setup. In practice the slowdown is usually smaller, often between ten and thirty percent, because a step also includes optimizer updates, communication and data loading. The exact number depends on the model, the hardware and how much you checkpoint.
How much memory does it save?
For deep networks, activation memory can fall from linear in depth to roughly the square root of depth when segments are chosen well. In real transformer training, people often see activation memory drop by a large factor, though total memory also includes weights and optimizer state, which are not reduced. Measure peak memory on your own model to know the real gain.
Is gradient checkpointing the same as activation checkpointing?
Yes, the two names refer to the same idea in most documentation. Some sources say activation recomputation or rematerialization. All of them describe saving fewer activations and recomputing the rest during the backward pass.
When should I not use gradient checkpointing?
Skip it when the model and batch already fit in memory with room to spare, when you only run inference, or when your main memory cost is parameters and optimizer state. In those cases it adds compute without solving the actual problem. Use sharding, lower precision or parameter efficient fine tuning instead.
Can I combine it with mixed precision and FlashAttention?
Yes, and this is the common setup. Mixed precision lowers the size of each activation, a fused attention kernel avoids storing the large attention matrix, and checkpointing discards the remaining inner activations. The three techniques address different parts of the memory bill and work well together.
Does it help with long sequences and large tables?
Yes. Activation memory grows with sequence length, and in table models it grows with the number of cells. Checkpointing is one of the main ways to train on long sequences or large tables without moving to a bigger accelerator.
What is the difference between checkpointing and gradient accumulation?
Gradient accumulation splits a batch into smaller micro batches and adds up their gradients, so only one micro batch of activations lives in memory at a time. Gradient checkpointing keeps the batch whole and discards inner activations, recomputing them later. The first costs more passes and the second costs more compute per pass, and they can be used together.
Does it work with distributed training?
Yes. Checkpointing happens inside each device, so it works with data parallel, tensor parallel, pipeline parallel and sharded strategies. Large training stacks usually include it as a configuration option.
Glossary
- Activation: the output of a layer during the forward pass.
- Activation checkpointing: another name for gradient checkpointing.
- Backward pass: the computation of gradients from the loss back to the inputs and parameters.
- Batch size: the number of examples processed together in one step.
- Checkpoint: in this context, a saved activation used as a restart point. Elsewhere, a saved copy of model weights.
- Forward pass: the computation from input to output and loss.
- Gradient: the derivative of the loss with respect to a parameter.
- Gradient accumulation: summing gradients over several micro batches before an update.
- Micro batch: a small slice of a batch processed in one pass.
- Offloading: moving tensors to slower memory to free fast memory.
- Peak memory: the highest memory use during a step.
- Recomputation: running part of the forward pass again to rebuild discarded values.
- Rematerialization: a compiler term for recomputation.
- Segment: a group of layers treated as one checkpointed unit.
- Sequence length: the number of tokens or cells the model reads at once.
- Sharding: splitting parameters, gradients or optimizer state across devices.
- Selective recomputation: recomputing only cheap, memory heavy operations.
- Throughput: tokens or samples processed per second.
- Reentrant checkpointing: an older implementation style that reruns the forward pass inside the backward call.
- Non reentrant checkpointing: a newer style that hooks into saved tensors and handles more cases.
Summary
Gradient checkpointing lowers training memory by storing only selected activations and recomputing the others during the backward pass. It leaves the math of training unchanged, costs roughly one extra forward pass per step, and can cut activation memory from linear to roughly square root growth in network depth. It works best when activations are the main memory cost, which is common for deep transformers, long sequences and large tables.
For tabular foundation models that read whole tables as context, it is a practical tool: it lets pre-training and fine tuning use larger tables, more rows and more columns within the same hardware limits. Use it together with mixed precision and fused attention, apply it at the block level, measure both memory and throughput, and spend the freed memory on larger batches or longer contexts. For more on how these models are built and evaluated, see /model/architecture, /model/benchmarks and /model/research.
Related
It belongs to Foundation models from scratch in my stack, under Foundation models and applied ML. Others in the same group: