Coming soonModel Training is launching soon in Scalata Enterprise. Teach a model to answer the way your team does.See what's coming

Resources/Model training/What happens during training

Concept

Gradient checkpointing

Gradient checkpointing saves memory by keeping only some of the intermediate results from the forward pass and recomputing the rest when the backward pass needs them. It uses much less memory and makes training about 20–30% slower.

Chapter 64 min read4 references

What it is

A training step has two halves. In the forward pass, each layer turns its input into an output, producing a set of intermediate results along the way: the activations. The backward pass then works through the layers in reverse to compute gradients, and each layer's gradients need that layer's activations. The simple approach keeps them all in memory until the backward pass reaches them. With long examples, activations can take more memory than the model's own weights.

Gradient checkpointing, also called activation checkpointing or recomputation, keeps only some activations, the checkpoints, and throws the rest away. When the backward pass needs a discarded one, it reruns the forward pass from the nearest checkpoint to recreate it. Chen et al. (2016) showed that this cuts activation memory for an n-layer network to the order of √n, for the price of about one extra forward pass per batch. In their experiments a 1,000-layer network went from 48 GB to 7 GB for about 30% more running time.

WITHOUT CHECKPOINTINGL1L2L3L4L5L6L7L88 sets of activations held until the backward passWITH CHECKPOINTINGL1L2L3L4L5L6L7L8recomputed from the last kept set when neededkept in memorydropped, recomputed
Eight layers, one training step. Normally every layer's activations stay in memory until the backward pass uses them. With checkpointing only a few are kept; when the backward pass reaches a gap, it reruns the forward pass from the nearest kept one. Less memory, roughly one extra forward pass of compute.

How it works in a language model

In transformer libraries the checkpoints usually sit at layer boundaries. Each layer's input is kept; everything computed inside the layer is thrown away and rebuilt when the backward pass gets there. So only one layer's full set of activations exists at any moment, instead of all of them.

Activation memory, kept and checkpointed (rough)
activations per GPT-style layer, 16-bit, efficient attention   (Korthikanti et al., 2022)
  ≈ 34 × batch × tokens × hidden size   bytes

illustrative 2B-class model: hidden size 2,048 · 28 layers · batch 2 × 1,024 tokens
  all kept                 34 × 2 × 1,024 × 2,048 × 28   ≈ 4.0 GB
  checkpointed, inputs      2 × 2 × 1,024 × 2,048 × 28   ≈ 0.23 GB
  + one layer rebuilt at a time                          ≈ 0.14 GB

How many checkpoints to keep is a dial. Keeping more means less to recompute but more memory; keeping fewer means the reverse. Checkpointing at every layer boundary is the usual default because it is simple and already removes most of the activation memory: what remains is one small input per layer plus the one layer being rebuilt.

The cost is compute. The forward pass effectively runs twice: once normally, and once piece by piece during the backward pass. The backward pass already costs about twice the forward pass, so the extra work is around a third of a step at most; in practice training runs about 20–30% slower. The result doesn't change: the recomputed values are the same ones, so the loss curve matches a run without checkpointing, up to tiny numerical differences.

Why large vocabularies make it matter more

At the very end of the model, every position in every example gets a score for every token in the vocabulary: the logits. Their size doesn't depend on the model's depth, only on how many positions are trained and how large the vocabulary is. Newer models have vocabularies of 150,000 to 250,000 tokens, so the logits alone can take gigabytes.

Memory for the logits
logits = batch × tokens × vocabulary size × bytes per number

2 × 1,024 × 248,000 × 4 bytes (32-bit)   ≈ 2.0 GB   before their gradients
A 2B model on a 16 GB GPU
  1. The model: about 2 billion parameters, from Qwen3.5's family, with a vocabulary of about 248,000 tokens.
  2. The logits: batch 2 × 1,024 tokens × 248,000 × 4 bytes ≈ 2 GB, and their gradients need about as much again.
  3. Everything else: the 16-bit weights, the adapter's training state and every layer's activations, all competing for what's left.
  4. Without checkpointing, a run like this can run out of memory in its first few steps. With it, the layer activations shrink to a fraction and the same run fits.
  5. Note: checkpointing doesn't shrink the logits themselves. It frees the room around them, which is what made the run fit.

Using it well

  • Reach for it before cutting data. Turning on checkpointing costs time; lowering max length cuts the ends off long examples. When memory is the limit, the slowdown is usually the better trade.
  • Spend the memory it frees. A larger batch or a longer max length can win back some of the lost speed per epoch.
  • Combine it with QLoRA for the largest models a machine can hold: QLoRA shrinks the weights, checkpointing shrinks the activations.
  • Know what it doesn't touch. Weights, gradients and optimizer state stay the same size, and it does nothing for running the finished model.
  • Check the default. Many fine-tuning recipes already switch it on. If memory is comfortable and speed matters, switching it off makes each step faster.

Common mistakes

  • Lowering max length first, cutting off the ends of long examples, when checkpointing alone would have made them fit.
  • Expecting it to shrink the weights or the optimizer state. It only affects activations.
  • Reading the 20–30% slowdown as a fault with the machine. It is the price of the memory saved.
  • Assuming it removes the logits' cost. With a huge vocabulary, a smaller batch or max length can still be needed.

References and further reading

External sources open in a new tab.

  1. 01Training Deep Nets with Sublinear Memory CostChen et al., 2016 · arXiv · arxiv.orgThe paper that made recomputation standard: √n memory for one extra forward pass.
  2. 02Reducing Activation Recomputation in Large Transformer ModelsKorthikanti et al., 2022 · arXiv · arxiv.orgWhere the per-layer activation estimate comes from, and how to recompute only the expensive parts.
  3. 03torch.utils.checkpointPyTorch docs · docs.pytorch.org
  4. 04GPU memory usageHugging Face Transformers docs · huggingface.coWeights, gradients, optimizer state and activations: where training memory goes.
Best practiceWill the model fit? Memory for trainingTraining needs room for the model's weights, its gradients and the optimizer's state, plus working memory. A simple rule of thumb, about 2 GB per billion parameters for LoRA and 12 GB for full fine-tuning, tells you before you start whether a model will fit your machine.MethodQLoRA (4-bit LoRA)QLoRA trains an ordinary LoRA adapter on a base model stored in 4 bits instead of 16. Training memory drops to roughly a third, so the same machine can fine-tune a model about three times larger, at some cost in speed.ConceptThe training loopFine-tuning is one short procedure repeated thousands of times: predict the target reply, measure how wrong the prediction was, work out which way each weight should move, and move it a little.SettingToken limits and trimmingEvery provider caps how many tokens one training example may contain. Long conversations over the cap are either trimmed — oldest exchanges removed first — or dropped.ConceptTokensThe chunks of text a model actually reads and writes — how they are made, why counts differ between models, and why every limit and price in training is measured in them.SettingHyperparametersHyperparameters are the training settings you choose, such as epochs, learning rate, batch size and adapter rank, as opposed to the weights the model learns. Auto picks them from your data and makes a good first run; this page explains each one so you can tune from there.

Turn weeks of expert work into minutes

See how Scalata fits your team, your data, and your controls.

Contact Sales