Why Training Is Its Own Beast

Training vs Inference gave the headline: training runs forward + backward + optimizer and needs ~16–20 bytes per parameter, against ~2 bytes for FP16 inference. This page makes that number concrete, because until you can size the training state you cannot answer the throughline for training: what does this cost in memory, and will it even fit on the GPUs I can rent?

What training has to keep in memory

Inference keeps one thing resident: the weights. Training keeps four:

• Weights — the parameters themselves.
• Gradients — one number per weight, telling you which way to nudge it. Same size as the weights.
• Optimizer state — Adam keeps a running mean and a running variance per weight (two extra full-size tensors), plus, in mixed precision, a high-precision FP32 master copy of the weights. So Adam alone is roughly the parameters again, twice over.
• Activations — the intermediate outputs from the forward pass, saved because the backward pass needs them to compute gradients. These scale with batch size and sequence length, not just parameter count.

The standard mixed-precision (FP16 + Adam) breakdown, per parameter:

  fp16 weights            2 bytes
  fp16 gradients          2 bytes
  fp32 master weights     4 bytes
  fp32 Adam momentum (m)  4 bytes
  fp32 Adam variance (v)  4 bytes
  ─────────────────────────────────
  persistent state       16 bytes / param   (+ activations on top)

That is the “~16 bytes/param” rule. Activations push the working figure toward ~16–20.

Worked example: training state for a 7B model

Take 7 billion parameters and apply the table:

weights      7e9 × 2 = 14 GB
gradients    7e9 × 2 = 14 GB
master copy  7e9 × 4 = 28 GB
Adam m       7e9 × 4 = 28 GB
Adam v       7e9 × 4 = 28 GB
─────────────────────────────
persistent  7e9 × 16 = 112 GB

112 GB before a single activation. Compare that to serving the same model: 14 GB of FP16 weights, which fits comfortably on one 80 GB GPU with room for a fat KV cache. The training state does not fit on that GPU at all — and adding activations for a real batch pushes it to 130–150 GB. This single fact is why training needs clusters: not because the math per token is exotic (the next page shows it is the same matmuls), but because the state outgrows any one card.

The optimizer is the hidden cost

The weights are 14 GB; the thing that blows the budget is the 98 GB of optimizer-and-gradient bookkeeping around them. When someone says “we can serve it but we can’t train it on this box,” this is what they mean.

Two levers that trade compute for memory

You usually cannot buy a 130 GB GPU, so training leans on two tricks that swap memory pressure for extra arithmetic.

Activation checkpointing (gradient checkpointing). Instead of storing every layer’s activations for the backward pass, store only a few checkpoints and recompute the rest on the fly during backward. Memory for activations drops from O(layers) to roughly O(√layers); the price is one extra partial forward pass, typically ~20–30% more compute. A textbook memory-for-compute trade — and on a memory-bound training job, very often worth it.

Gradient accumulation. Large batches stabilize training, but a large batch means large activations, which may not fit. So split the batch into micro-batches, run forward+backward on each, add up the gradients, and only step the optimizer after K micro-batches. You get the statistics of a batch K times larger than fits in memory, paying only in wall-clock time (you ran K sequential micro-batches instead of one big parallel one).

for micro_batch in split(batch, K):
    loss = forward(micro_batch)
    backward(loss)          # gradients accumulate, not reset
optimizer.step()            # one update for the whole effective batch
optimizer.zero_grad()

Both levers point the same direction: when memory is the wall (and in training it usually is), spend compute to climb it.

Forward + backward, in one breath

Why is backward ~2× the FLOPs of forward? Each weight participates in two gradient computations — the gradient with respect to its input (to keep propagating backward) and the gradient with respect to the weight itself (to update it). So a training step is ~6× params FLOPs per token versus ~2× for inference, and it carries all that state alongside. More arithmetic, far more memory, paid once — but “once” can still mean thousands of GPU-hours.

It actually happened — sublinear-memory training (Chen et al., 2016)

The O(√layers) figure in the activation-checkpointing section isn’t a rule of thumb — it’s a specific result. In Training Deep Nets with Sublinear Memory Cost (Tianqi Chen and colleagues, 2016), the authors showed that by storing only a handful of evenly spaced activation checkpoints and recomputing the rest during the backward pass, the activation memory of an n-layer network drops from O(n) to roughly O(√n) — at the cost of about one extra forward pass, the ~20–30% compute overhead noted above. It is one of the cleanest memory-for-compute trades in deep learning, and it predates the LLM era: the technique was built for vision and RNN training and became indispensable once transformer activations started dwarfing the weights. Every time you fit a bigger batch or a longer sequence by toggling “gradient checkpointing” in a training config, this is the 2016 result doing the work.

Check your understanding

1. List the four categories of data training must hold in memory, and which one inference also needs.
2. Using the ~16 bytes/param rule, compute the persistent training state for a 3B model.
3. Why can you often serve a model on one GPU but not train it there?
4. What does activation checkpointing trade away, and what does it buy?
5. A team’s batch won’t fit in memory but they need its statistics for stable training. Which technique solves this, and what’s the cost?

Show answers

1. Weights, gradients, optimizer state (e.g. Adam’s momentum, variance, and a master copy), and saved activations. Inference needs only the weights (plus a runtime KV cache).
2. 3e9 × 16 bytes = 48 GB of persistent state, before activations.
3. Serving needs ~2 bytes/param (just weights); training needs ~16 bytes/param because of gradients, optimizer moments, and the master copy — about 8× more memory, which overflows a single card.
4. It trades extra compute (recomputing activations during the backward pass, ~20–30% more) for a large reduction in activation memory (roughly O(layers) → O(√layers)).
5. Gradient accumulation — run several micro-batches, sum their gradients, and step once. The cost is wall-clock time, since the micro-batches run sequentially rather than as one large parallel batch.