Why training needs so much more GPU memory than inference

During inference, the model performs a forward pass and generates outputs. During training it must later reverse that computation and determine how every trainable weight should change.

That means training needs to preserve intermediate information from the forward pass.

Suppose:

During inference, once \(y\) has been consumed, \(x\) may no longer need to remain in memory.

During training, the weight gradient is:

so the backward pass needs \(x\).

Scale that across dozens of layers and thousands of tokens: saved activations become enormous.

Four major memory consumers¶

A simplified training-memory budget is:

For a 70B model:

Weights¶

At BF16:

Gradients¶

If BF16:

Adam optimizer state¶

Adam tracks \(m\) and \(v\). If FP32:

for each state, or:

together.

Even before activations, the crude total is already ~840 GB.

Real systems vary in precision and representation, but the scale is the point.

Activations¶

Suppose a layer has hidden states:

with \(B=8\), \(T=8192\).

That is:

values.

At BF16:

for just one large activation tensor.

Transformer blocks create multiple such intermediates.

Activation checkpointing¶

Instead of storing every activation:

L1 → A1
L2 → A2
L3 → A3
...

save only selected checkpoints. During backward, recompute missing activations from the nearest checkpoint.

This deliberately trades more FLOPs for less memory:

Why multi-GPU training is unavoidable¶

A 70B BF16 model’s weights alone are ~140 GB. Training state is much larger.

Hence: - pipeline parallelism splits layers, - tensor parallelism splits matrix operations, - data parallelism splits examples, - FSDP/ZeRO-like techniques shard weights, gradients, and optimizer state.

Once the model spans many accelerators, communication becomes part of the algorithm.

Training also stores temporary computation¶

The static parameter/gradient/optimizer accounting is only the beginning. During a forward pass through a Transformer block, training may need information related to:

• residual-stream activations,
• normalization inputs,
• Q/K/V projections,
• attention intermediates,
• MLP inputs and outputs,
• dropout masks or other state.

Saving all of those naively across long sequences can dominate memory.

Activation checkpointing says: save a smaller set of checkpoints and recompute the rest later. If a forward operation takes 1 unit of compute and its activation takes hundreds of MB, recomputation can be the economically sensible choice.

Why inference has a different memory problem¶

Inference avoids gradients and optimizer states, but autoregressive inference introduces the KV cache.

So the dominant memory components differ:

Training:
weights + gradients + optimizer + activations

Inference:
weights + KV cache + runtime buffers

That is why “this model fits for inference” does not imply “this model can be trained on the same GPUs.”

Sharding changes ownership, not total information¶

Techniques such as FSDP or ZeRO do not magically eliminate optimizer state. They distribute it.

If 560 GB of Adam moments are sharded evenly over 64 workers, each worker owns roughly:

instead of 560 GB.

But now training depends on moving the right parameter/state shards to the right place at the right time.

Memory savings become communication costs.

Activation memory grows with sequence length and batch size¶

Parameter memory is mostly fixed for a given model, but activation memory changes with the training workload.

If one saved hidden-state tensor has shape:

then doubling batch size doubles that tensor, and doubling sequence length doubles it again.

Attention-related intermediates can scale even more sharply with \(T\), which is why long-context training is particularly memory-intensive.

This creates an important systems distinction:

Training a model at 2K context and training the same model at 128K context can require radically different memory strategies even though the parameter count is identical.

Sequence parallelism, context parallelism, FlashAttention, and activation recomputation all become more valuable as context grows.

Why optimizer precision can differ from model precision¶

Even when the forward weights are BF16 or lower precision, the optimizer may maintain higher-precision master values or statistics because updates can be extremely small.

This is another recurring mixed-precision principle:

The forward pass, gradients, optimizer moments, and communication buffers need not all use the same numeric format.