LLM Memory Math

The four buckets

Training memory splits cleanly into four categories, and almost everything you read about fitting models on hardware is really a statement about one of them:

1. Parameters — the weights themselves, the thing you are actually learning. 2. Gradients — one number per parameter, produced by the backward pass. 3. Optimizer state — the bookkeeping an optimizer like Adam keeps per parameter so it can adapt each step, plus a high-precision master copy of the weights. 4. Activations — the intermediate tensors from the forward pass that the backward pass needs in order to compute gradients.

The first three scale only with the parameter count: they are fixed the moment you choose a model and an optimizer, independent of how much data you feed. The fourth scales with the workload — batch size, sequence length, and layer count — and is the one you tune at run time. Get these four numbers and you can predict, before launching anything, whether a run will fit.

Parameters: the base weight

A parameter is just a number, and its memory cost is its numeric width. In old-school full precision every weight is an fp32 float — 4 bytes. Modern training keeps a working copy in a 16-bit format, bf16 or fp16 — 2 bytes. So a model with P parameters needs 2P bytes for a bf16 weight copy, 4P for fp32.

The arithmetic is trivial and worth internalising because it anchors everything else. A 7B model is 7 × 10^9 parameters, so a bf16 copy is 7e9 × 2 = 14 GB; a 70B model is 140 GB in bf16 — already larger than a single 80 GB accelerator before you add a single gradient. This is the number people quote for inference, and it is the reason the weights feel expensive. In training they are, counter-intuitively, the cheapest of the four buckets.

Gradients: one number per parameter

The backward pass computes ∂L/∂w for every weight w, so gradients are the same shape as the parameters: exactly one value per parameter. In a bf16 training loop the gradients are typically stored in 2 bytes each as well, giving another 2P bytes; some recipes accumulate them in fp32 (4P) for numerical stability.

So after parameters plus gradients you are already at roughly 4P bytes (both bf16) — a 7B model is 28 GB for weights and gradients together. There is no trick that removes gradients entirely while you are doing gradient descent; they are the signal. What you can do is avoid holding all of them at once (sharding), or avoid computing them for most weights at all (parameter-efficient fine-tuning, where only a small adapter has gradients). Keep that distinction in mind — it is the seam several memory-saving methods exploit.

Optimizer state: where Adam gets expensive

A plain SGD optimizer keeps nothing extra: it nudges each weight by its gradient. But almost nobody trains large models with plain SGD. The workhorse is Adam (or AdamW), which adapts the step size per parameter using two running statistics: a first moment m (an exponential moving average of the gradient) and a second moment v (an EMA of the squared gradient).

Both m and v are one value per parameter, and both are kept in fp32 for stability — that is 4 + 4 = 8 bytes per parameter on top of the weights and gradients. Adam also holds an fp32 master copy of the weights (another 4 bytes), because the tiny updates would vanish if applied directly to the 16-bit working copy. That master copy plus the two moments — 4 + 4 + 4 = 12 bytes per parameter — is the real reason training memory dwarfs inference memory. The optimizer, not the model, is the heavyweight.

The 16 to 18 bytes per parameter rule

Now add the standard mixed-precision Adam buckets together, per parameter:

bf16 weight copy      2 bytes
bf16 gradient         2 bytes
fp32 master weight    4 bytes   (optimizer)
fp32 Adam moment m    4 bytes   (optimizer)
fp32 Adam moment v    4 bytes   (optimizer)
---------------------------------
TOTAL              = 16 bytes / parameter

That is the famous 16 bytes per parameter figure for mixed-precision Adam. If you accumulate gradients in fp32 instead of bf16 the gradient line becomes 4 bytes and the total rises to 18 bytes/param. (Curiously, full-fp32 Adam — 4 + 4 + 4 + 4 — also lands at 16 bytes, so the rule is robust.) Multiply by the parameter count and you have the fixed, workload-independent core of your budget: 16P. A 7B model needs 7e9 × 16 = 112 GB for parameters, gradients, and optimizer state alone — before a single activation.

Activations: the bucket that scales with your batch

The fourth bucket is different in kind: it does not scale with the parameter count but with the work. To compute gradients, autograd must remember the forward-pass intermediates — layer inputs, attention scores, the outputs feeding each nonlinearity. Roughly, activation memory grows as batch × sequence × hidden × layers, and the attention scores add a term that grows with sequence^2.

A useful approximation (Korthikanti et al.) puts per-layer activation memory at about s·b·h × (34 + 5·a·s/h) bytes in 16-bit, where s is sequence length, b batch, h hidden size, a attention heads. The key property is that this is linear in batch and layers and roughly quadratic in sequence. Unlike the first three buckets, activations are entirely under your control at run time: halve the batch, halve this bucket. That is exactly why activations are where most practical memory tuning happens.

A worked breakdown: a 7B model on one accelerator

Put it together for a 7B model (h = 4096, L = 32 layers, a = 32 heads) trained with mixed-precision AdamW at sequence length s = 2048, batch b = 1 per device:

Parameters (bf16)          14 GB
Gradients (bf16)           14 GB
Optimizer state (fp32)     84 GB   (12 bytes/param)
  = 16 bytes/param total  112 GB
Activations (no ckpt)     ~30 GB   (grows with b and s)
---------------------------------------
TOTAL                    ~142 GB

142 GB does not fit on an 80 GB accelerator — not even close, and this is with a batch of one. That single sum explains why 7B is about the ceiling for naive single-GPU full fine-tuning, why the optimizer state (84 GB, the largest slice) is the first thing sharding attacks, and why every larger model is a multi-device or memory-offload story from the outset.

What each knob actually reduces

Every memory-saving technique is best understood by which bucket it drains. Confusion comes from treating them as interchangeable ‘make it smaller’ switches; they are not, and stacking two that hit the same bucket gives less than you expect.

Read the table as a diagnosis tool: first identify which bucket is over budget, then pick a knob that drains that bucket. An out-of-memory error at batch 1 is an optimizer-state or parameter problem — reach for sharding, offload, or PEFT, not a smaller batch.

Sharding: the memory you do not have to hold

Sharding (ZeRO in DeepSpeed, FSDP in PyTorch) attacks the biggest fixed buckets by refusing to store the whole of them on any one device. With N data-parallel workers, ZeRO stage 1 splits the optimizer state N ways, stage 2 adds the gradients, and stage 3 shards the parameters themselves. Each device holds only its 1/N slice and gathers the rest on demand during the step.

The effect on our 7B example is dramatic: at stage 3 across 8 devices the 112 GB of fixed state becomes about 14 GB per device, which now fits comfortably with room for activations. The cost is communication — devices must exchange shards every step — so sharding trades interconnect bandwidth for memory. This is why the same 16-bytes-per-parameter arithmetic that condemns single-GPU training also tells you exactly how many devices a run needs: divide the fixed budget by the memory you can spare per device.

CPU and small-model implications

On a CPU box or a modest edge device the four-bucket arithmetic is the difference between a plan and a stall. Full fine-tuning of even a 1-3B model implies 16-48 GB of fixed state, which is why CPU-side training almost always means PEFT: freeze the base weights (a single 2-byte, or with 4-bit quantization sub-1-byte, copy and no gradient or optimizer state for them) and train a small LoRA adapter whose parameters number in the millions, not billions.

That collapses buckets two and three to near nothing — QLoRA fine-tunes models that would otherwise need a data-centre’s worth of accelerators on a single card — because the expensive optimizer state now rides on a few million adapter weights rather than every weight in the model. For pure inference on CPU the story is simpler still: only the parameter bucket exists (plus a small KV cache), so a 4-bit 7B model in roughly 4-5 GB is entirely tractable.

Pitfalls that break the accounting

The envelope estimate is reliable, but a few real-world costs sit outside the four clean buckets and routinely cause an out-of-memory error that ‘should not’ happen. Memory fragmentation: the allocator may fail to place a tensor even when the free total is sufficient, so usable memory is always somewhat below the nominal figure. Peak versus steady state: transient spikes — a large logits tensor over the vocabulary, gradient all-gathers, temporary buffers — set the ceiling, not the average, so a run can crash at the peak of step one.

Also easy to forget: the CUDA context and framework overhead claim a fixed slice before your model loads; a long-context run pays a KV cache that scales with sequence length; and gradient accumulation raises the effective batch without raising activation memory, which is the cheapest way to simulate a big batch you cannot afford to hold. Budget for peak, not average, and leave headroom — the accounting tells you the floor, never the exact ceiling.