How a model too big for any GPU is trained on thousands of them, with every byte of memory, every message on the network, and every idle second counted.

Here is a puzzle that has ended many first attempts at training a large model. A 7-billion-parameter model stored in bfloat16 takes 13.5 GB. An H100 has 80 GB of memory, enough to hold that model almost six times over. So you load the model onto one GPU, run a forward pass, call loss.backward(), call optimizer.step(), and the program dies with CUDA out of memory. Where did the other 66 GB go?

The answer is that the weights are the smallest of four bills. Training also stores a gradient for every weight, an optimizer's private state for every weight, and the intermediate results of the forward pass that the backward pass needs, and each of these can be larger than the weights themselves. This section counts every byte. By the end you will be able to look at a model's shape and a batch size and say, to within a few percent, how much memory training it takes, which is the first question behind every decision in the rest of this chapter.

Four kinds of memory

Take a model with $N$ parameters, trained with Adam on micro-batches of $b$ sequences of $T$ tokens each. The GPU holds four things.

1. Parameters, the weights themselves, $N$ numbers.
2. Gradients, $\partial\mathcal{L}/\partial\theta$, one number per parameter, produced by the backward pass and consumed by the optimizer.
3. Optimizer state. Adam keeps a running mean $\mathbf{m}$ and a running mean of squares $\mathbf{v}$ of the gradients (§0.6), two more numbers per parameter. Mixed-precision training adds a full-precision copy of the weights, the master weights, which we will justify in a moment.
4. Activations, the tensors computed in the forward pass and saved for the backward pass. These scale with $b$ and $T$, not with $N$.

Rajbhandari et al. (2020), the ZeRO paper we meet in §10.3, call the first three model states and everything else residual states (activations plus temporary buffers and fragmentation). The split is useful because the two halves scale differently. Model states grow with the size of the model and are the same for every batch. Activations grow with the batch and the sequence length and can be traded for compute.

Why must activations be saved at all? Recall from §0.3 that the backward pass of $\mathbf{Y} = \mathbf{X}\mathbf{W}$ needs $\partial\mathcal{L}/\partial\mathbf{W} = \mathbf{X}^\top\mathbf{G}$, where $\mathbf{G} = \partial\mathcal{L}/\partial\mathbf{Y}$ is the upstream gradient. The input $\mathbf{X}$ was computed during the forward pass, long before $\mathbf{G}$ exists, so the framework must keep it until the backward pass reaches this layer. Every matrix multiply keeps its input, every nonlinearity keeps its input or output, and every dropout keeps its mask. For a deep network, that is a lot of tensors held at once, because the forward pass of the last layer finishes before the backward pass of the first layer begins.

Adam's bill, 16 bytes per parameter

Start with plain float32 training. Each number takes 4 bytes, and there are four numbers per parameter (weight, gradient, $\mathbf{m}$, $\mathbf{v}$), so

Now mixed precision, the way every large model is trained. The forward and backward passes run in a 16-bit format (bfloat16 today, float16 in older recipes), because tensor cores are several times faster in 16 bits and the activations take half the space (§8.5 covers the formats in depth). So the working copy of the weights is 16-bit, 2 bytes, and so are the gradients the backward pass produces, another 2 bytes. The optimizer, however, keeps everything in float32. It holds $\mathbf{m}$ and $\mathbf{v}$ (4 bytes each) and a float32 master copy of the weights (4 bytes). Each step it updates the master weights in float32 and then casts them to 16 bits for the next forward pass. The total is

In plain words, mixed precision does not save any memory on model states. It moves 4 bytes from the weights and gradients into a master copy, and its real savings are in activations and speed.

ZeRO writes this as $2N + 2N + KN$ with $K = 12$ bytes of optimizer state per parameter, and you will see "$16N$" or "$16\Psi$" (the paper's symbol for the parameter count) all over the literature. For a 7B model that is 112 GB, which already exceeds the 80 GB of an H100 before a single activation is stored. That answers the puzzle.

Why keep a float32 master copy at all? Because bfloat16 cannot represent small changes to a number. It has 8 bits of significand precision (7 stored plus an implicit leading 1), so the gap between 1 and the next representable number is $2^{-7} \approx 0.0078$, and an update smaller than half the local gap, about $0.002$ to $0.004$ for a weight near 1 (roughly $1/256$ of the weight), rounds away to nothing. Adam's update for one weight is roughly the learning rate times a number of order 1, and learning rates for large models are around $10^{-4}$, so a weight kept only in bfloat16 would never move. A float32 master copy accumulates every one of those small steps, and its bfloat16 cast follows it. §8.5 works this out bit by bit, including the experiment where a thousand updates leave a pure bfloat16 weight exactly where it started, and the alternatives (stochastic rounding, compensated summation) that let some recipes drop the copy. For the memory bill, what matters is that the standard recipe needs the master weights, so it stays at 16 bytes.

Several variants are common, and an interviewer will expect you to know their costs.

The third row is what Megatron-LM does for bfloat16 training, accumulating gradients over micro-batches in float32 so that summing many small contributions does not lose precision. The fourth uses the block-wise quantized optimizer states of Dettmers et al. (2022). The fifth only works with tricks such as stochastic rounding or compensated (Kahan) summation that recover the lost low bits, and it is used mostly for fine-tuning. Optimizers with less state, such as Adafactor's factored second moment or Muon, are the subject of §8.6.

def model_state_bytes(N, recipe):
    total = 16 * N
    return {"weights": 2 * N, "grads": 2 * N, "master": 0, "optimizer": 0, "total": total}

Activations, counted one tensor at a time

Model states are the easy half. Activations are where people's estimates go wrong, because they depend on exactly which tensors each operation saves. The standard reference is Korthikanti et al. (2022), who counted them for a GPT-3-style layer. We will redo their count from scratch, then check it against PyTorch.

Set up the layer. The input is a tensor of shape $(b, T, d)$, a micro-batch of $b$ sequences of $T$ tokens of width $d$, and the attention has $h$ heads. The layer is pre-norm GPT, with LayerNorm, attention, an output projection, dropout, a residual add, then LayerNorm, a $d \to 4d$ linear, GELU, a $4d \to d$ linear, dropout, and another residual add (the block of §5.6, with dropout after the attention weights and after each sublayer). Activations are stored in 16 bits (2 bytes per number) and dropout masks in 1 byte per number. Write $Tbd$ for the number of entries of one $(b, T, d)$ tensor. Here is every tensor the backward pass needs.

Add the columns. The attention block stores $2 + 4 + 2 + 2 + 1 = 11$ units of $Tbd$ plus $5hT^2b$, the MLP $2 + 8 + 8 + 1 = 19$ units, and the two norms 4 units. The total for one layer is

In plain words, each token costs about $34d$ bytes per layer for everything linear, plus $5hT$ bytes for its row of every head's attention matrix, and the second term grows with the sequence length.

This is Korthikanti et al.'s equation (1), written in our symbols (they write $s$ for the sequence length, $h$ for the width, and $a$ for the heads, so their $sbh(34 + 5as/h)$ is our formula). The LayerNorm also saves a mean and a variance per token, $2Tb$ numbers, which they drop because $d$ is in the thousands.

Notice how lopsided the two terms are. For GPT-3 ($d = 12{,}288$, $h = 96$, $T = 2048$), $5hT/d = 80$, so the attention matrices are 70% of all activation memory. For a 7B model at $T = 4096$ ($d = 4096$, $h = 32$), $5hT/d = 160$, more than four times the 34 of everything else. The quadratic term is the score matrix of §4.7, and it is exactly what FlashAttention never stores (§12.3). Korthikanti et al. reach the same place from the other direction with selective recomputation, which drops these tensors and recomputes them in the backward pass. Either way the per-layer cost falls to $34\,Tbd$.

The bluntest version of the same trade is activation checkpointing, also called gradient checkpointing or full recomputation (Chen et al. 2016). Keep only each layer's input, $2Tbd$ bytes, throw away everything the layer computed, and when the backward pass reaches that layer, rerun its forward pass from the saved input to rebuild the activations just in time. Memory per layer falls from $34\,Tbd$ (or more) to $2Tbd$, and the price is one extra forward pass per step, about a third more compute, which §10.7 counts exactly. The three modes, none, selective, and full, are the three settings of the exercises and the widget below.

PyTorch lets you watch the saving happen. torch.autograd.graph.saved_tensors_hooks calls a function every time an operation saves a tensor for backward. We record each saved storage once (several views can share one storage), skip the parameters, and charge 2 bytes per floating-point entry and 1 byte per boolean mask, as a bfloat16 run on a GPU would. Dropout is written out with an explicit boolean mask, which is what the fused GPU kernel stores.

The measured total is the formula plus 1,024 bytes, and those 1,024 bytes are exactly the LayerNorm statistics. Every other byte matches the table.

The whole model, and the modern layer

Multiply by the number of layers $L$ and you have nearly all the activation memory. Korthikanti et al. add the pieces outside the layers. The embedding dropout mask and the final norm are small, but the logits are not. The output layer produces $bT \times V$ logits, and computing the loss in float32 keeps $4bTV$ bytes. For Llama 3's vocabulary of $V = 128{,}256$ at $T = 8192$, that is 4.2 GB per sequence, as much as several layers. Fused cross-entropy kernels that process the vocabulary in chunks avoid ever materializing it.

The 34 is specific to the GPT-3 layer. A Llama-style layer has no dropout, RMSNorm instead of LayerNorm, grouped-query attention with $n_{kv}$ key-value heads of width $d_k$, a SwiGLU feed-forward network of width $d_{\text{ff}}$ (§5.3), and FlashAttention. The same bookkeeping, done for plain PyTorch operations, gives per token and per layer the two norm inputs and two norm outputs ($8d$ bytes), $\mathbf{Q}$ and the attention output ($4d$), $\mathbf{K}$ and $\mathbf{V}$ ($4n_{kv}d_k$), FlashAttention's per-head log-sum-exp in float32 ($4h$), and four $d_{\text{ff}}$-wide tensors in the SwiGLU (the gate, the up projection, the activated gate, and their product, $8d_{\text{ff}}$).

The Llama layer comes to about $41d$ bytes per token, a bit more than GPT-3's $34d$ because its feed-forward network is wider. Fused kernels that recompute the activated gate shave a few $d$ off. The lesson is not the exact constant, which depends on your implementation, but the method. List the tensors each operation saves, multiply by their sizes, and check the result against the framework, as we did above.

def layer_activation_bytes(T, b, d, h, recompute="none"):
    return 34 * T * b * d + 5 * h * T * b          # placeholder: the attention term is wrong

Worked example, a 7B model on one GPU

Put the pieces together for Llama 2 7B ($L = 32$, $d = 4096$, $h = 32$, $N = 6.74$ billion) at $T = 4096$, one sequence per micro-batch, and see how much each choice buys. We use the GPT-3 constant 34 for the layers, which is close for this model, and write GB for $10^9$ bytes.

Even with every activation trick, training this model with plain Adam needs about 109 GB, more than an 80 GB GPU holds, and nearly all of it is model states. No amount of recomputation helps, because the activations are already small. The fixes all attack the model states. You can shard them across several GPUs (ZeRO and FSDP, §10.3), shrink them (8-bit optimizer states, or freezing the weights and training a small adapter as in LoRA, which keeps only the frozen 16-bit weights plus a tiny optimizer), or move them to CPU memory.

Without FlashAttention or selective recomputation the picture is different. The activations alone are about 105 GB, and they grow by that much again with every extra sequence in the micro-batch. Which term dominates depends on $T$, $b$, and the model, so the habit to build is to compute both.

Worked example, a 70B model

Now Llama 3 70B ($L = 80$, $d = 8192$, $h = 64$, $N = 70.6$ billion) at $T = 8192$. Model states are $16 \times 70.6 = 1{,}129$ GB, so even if the states were spread perfectly over 80 GB GPUs you would need 15 of them before storing any activations. One sequence's layer activations are 1.9 TB with naive attention, 182 GB with FlashAttention, and 10.7 GB with full recomputation, and the float32 logits over Llama 3's 128,256-token vocabulary add another 4.2 GB on top of each.

So no trick fits a 70B model on one GPU, and even a whole node of eight 80 GB GPUs, with 640 GB, holds only about half of its model states. Training it takes many nodes and a way to split the model states across them (short of offloading them to CPU memory). The rest of this chapter is about the ways to split, what each one costs in communication, and how to combine them. The widget below lets you see the bill for several models and the effect of each lever, including the ZeRO sharding of §10.3.

def max_micro_batch(N, L, d, h, V, T, mem_bytes, recipe="mixed", recompute="none"):
    per_seq = L * 34 * T * d
    return int((mem_bytes - 16 * N) // per_seq)

The memory you do not see

A calculation like the one above gets you within about 10% of the truth, and the remaining 10% has three sources worth knowing.

Temporary buffers. Some operations need workspace while they run, for example a matrix multiply's scratch space, the full-precision copy of a gradient during clipping, or the flattened buffers that communication libraries use to send many small tensors at once. They come and go within a step, but the peak is what matters.

Fragmentation. PyTorch's caching allocator keeps freed blocks to reuse them, and when tensors of many different sizes are allocated and freed, the free memory can be broken into pieces too small to serve the next large request. You can then see an out-of-memory error while torch.cuda.memory_reserved() is well above torch.cuda.memory_allocated(). Setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True lets the allocator grow segments instead of creating new fixed-size ones, which often helps.

The peak is not at the end of the forward pass. During the backward pass, activations are being freed while gradients are being created, and the optimizer step can create temporaries the size of the model states. When memory is tight, measure.

The snapshot shows every live tensor over time, colored by the line of code that allocated it, and it is the fastest way to find out which part of a step sets the peak.

Read more

• Rajbhandari et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. Section 3 is the clearest account of where training memory goes, and the source of the $2 + 2 + 12$ accounting.
• Korthikanti et al. (2022). Reducing Activation Recomputation in Large Transformer Models. Section 4 derives the $34 + 5as/h$ formula tensor by tensor, with and without tensor and sequence parallelism.
• Micikevicius et al. (2018). Mixed Precision Training. Master weights, loss scaling, and why 16-bit training needs a 32-bit copy.
• Dettmers et al. (2022). 8-bit Optimizers via Block-wise Quantization. How to keep Adam's state in one byte per number without losing accuracy.
• Chen et al. (2016). Training Deep Nets with Sublinear Memory Cost. The original activation checkpointing paper, trading compute for memory.
• Tazi et al. (2025). The Ultra-Scale Playbook. Its first chapter profiles the memory of a real training step and is a good companion to this section.
• PyTorch documentation. Understanding CUDA Memory Usage. How to record and read memory snapshots.

Watch next

▶

▶