Where the Memory Goes

An 8-billion-parameter model in bf16 is 16 GB of weights. An H100 has 80 GB. So you can train four of them on one GPU, and the whole industry's obsession with sharding is some kind of collective delusion.

Obviously not. The actual answer is that you cannot train one of them on one GPU — not even close, not even at batch size 1. Getting from "16 GB of weights" to "doesn't fit in 80 GB" is the subject of this chapter, and it's worth doing carefully, because the sharding techniques in chapters 7 through 11 are all attacks on specific terms in this equation. If you don't know which term is biggest, you can't know which technique to reach for.

There are exactly four consumers of memory during training:

The split matters more than the four terms individually. The first three are fixed — set by parameter count and your choice of optimizer, unchanging across the run. The last is elastic: you control it directly through batch size and sequence length, and it can be made almost arbitrarily small at the cost of FLOPs. Different problems, different solutions.

The fixed cost: 16 bytes per parameter#

Let's count what mixed-precision training with Adam actually keeps resident. The word "mixed" is doing real work here — the reason this is expensive is that "training in bf16" does not mean everything is in bf16.

Two of these surprise people.

The master weights. Why keep a second, fp32 copy of every weight when you already have a bf16 one? Because bf16 has 8 bits of mantissa, giving roughly 3 significant decimal digits. A typical update is on the order of $\eta \cdot \hat{m} \approx 10^{-3} \times$ the weight's own magnitude. Adding a number to another number a thousand times larger, in a format with three digits of precision, rounds to a no-op. The update silently vanishes and training stalls — not diverges, stalls, which is much harder to diagnose. So the optimizer step happens in fp32 against the master copy, and the bf16 weights are re-derived from it after each step.Micikevicius et al., Mixed Precision Training, ICLR 2018. The paper is about fp16 and spends much of its length on loss scaling; bf16 later made that part unnecessary, but the master-weights argument survives intact.

The optimizer state is the single largest term. Adam's two moments are 8 bytes per parameter — four times the bf16 weights, and half of the entire fixed budget. This is the direct reason ZeRO exists, and the reason its stage 1 (shard the optimizer state, nothing else) already gets you most of the way there.

Run the numbers on our anchor:

An 8B model exceeds a single H100 by 1.5× before storing one activation. This is the reason distributed training isn't optional above toy scale — not throughput, not wall-clock. It doesn't fit.

The elastic cost: activations#

Now the term that actually varies, and the one that generates the interesting engineering.

During the forward pass, every intermediate tensor needed to compute a gradient in the backward pass has to be kept. Not just layer outputs — the input to every matmul, the output of every nonlinearity, every dropout mask. The autograd graph pins all of it until backward consumes it.

Korthikanti et al. counted this precisely for a standard transformer block. Per layer, in bf16:Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, MLSys 2023. Their §2 derivation is the canonical reference and worth reading in full; the constant 34 is architecture-specific but the two-term structure is universal.

where $s$ is sequence length, $b$ micro-batch size, $h$ the hidden dimension, and $a$ the number of attention heads.

Don't let the constant distract you. The structure is the point, and there are two terms:

• $34sbh$ — linear in sequence length. Every residual-stream-sized tensor in the block: layernorm inputs and outputs, Q/K/V, the projection input, the MLP's intermediates, dropout masks. About seventeen tensors' worth.
• $5as^2b$ — quadratic in sequence length. This is the attention score matrix and its friends: $QK^\top$, the softmax output, the dropout mask over it. One $s \times s$ matrix per head.

Those two terms behave completely differently, and essentially every technique in the next two chapters targets one or the other. Concretely, for Llama 3 8B at $s = 8192$, $b = 1$, summed over all 32 layers:

At batch size one. The quadratic term is nine times everything else combined, and 354 GiB against an 80 GiB GPU means a naive implementation of Llama 3 8B cannot do a single forward-backward pass at its own context length.

Which is a good moment to notice the year. This calculation is why long context was genuinely hard before 2022, and why two independent fixes arrived at roughly the same time: flash attention deletes the quadratic term by never materializing the score matrix in HBM (chapter 4), and selective recomputation deletes it by throwing it away and recomputing it (chapter 3). They arrive at the same place from opposite directions.

Kill the quadratic term and you're at 34 GiB. Add full recomputation on top — keep only each layer's input, 2sbh — and you're at 2 GiB. That's a 177× range on the same model at the same batch size, entirely determined by which activation strategy you picked.

The calculator#

Every number above comes from this, and it's the backbone widget for the rest of the book — chapters 3, 7, 8, and 11 all come back to it. Start with the defaults, then turn flash attention off and watch what 2021 felt like.

Some things worth doing with it:

• Set the model to 70B with ZeRO off. Nothing you can do to batch size or sequence length saves you; the fixed term alone is 13× an H100. Sharding isn't an optimization here, it's the price of entry.
• Set ZeRO stage 3 with 64 GPUs on the 8B model. The fixed term nearly vanishes and activations become essentially the whole bill — which is why ZeRO and activation checkpointing are complements, not alternatives.
• Push sequence length to 32k with flash off. The quadratic term goes superlinear in a way that no amount of sharding fixes, because it's per-device and per-layer.

The terms nobody counts#

The four terms above are the honest core, but they're not the whole allocation. In a real run you should expect to lose another 10–20% to:

Fragmentation. PyTorch's caching allocator grabs large blocks from CUDA and sub-allocates. Varying sequence lengths across micro-batches produce oddly-sized requests that don't fit reclaimed blocks, so reserved memory drifts above allocated memory. When nvidia-smi shows 78 GB used and torch.cuda.memory_allocated() shows 62 GB, that gap is fragmentation. PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True helps materially.

Communication buffers. DDP's gradient buckets, ZeRO's all-gather staging areas, NCCL's internal buffers. Typically single-digit GB, and they scale with world size and bucket configuration.

The logits. Easy to forget and occasionally enormous: the output tensor is $s \times b \times V$, and with $V = 128256$ that's 1.96 GiB per 8k sequence in bf16 — and the cross-entropy computation usually wants it in fp32, doubling that, plus a same-sized gradient. On small models with large vocabularies this can rival a transformer layer. Chunked loss computation exists for exactly this reason.

CUDA context and kernels. 0.5–1 GiB per process before you allocate anything.

What this chapter buys you#

You now know which term to attack, which is the only question that matters:

The next chapter takes the largest controllable term — activations — and spends FLOPs to make it disappear. It's the purest example of the three-budget trade in the book, and the source notes I'm building on give it two sentences.