Why training needs so much memory.

Running a model only needs its weights, the millions or billions of numbers it learned. Teaching it needs those weights plus a pile of notes about every one of them: which way each number should move, how it has been moving lately, a more precise backup copy, and everything the model worked out on the way to its answer so it can trace its mistakes backwards. Add up the notes and they come to many times the weights.

This is the arithmetic that decides what you can train on your own machine. Get it wrong and a job you were sure would fit dies before it finishes a single step, or you rent cloud GPUs for something the Spark could have done. Get it right and you can say, before you run anything, which model sizes are yours to train, which need tricks, and which need a cluster. Interviewers expect you to do this sum on a whiteboard.

It builds on How much memory does a model need?, which sized the serving side, and borrows the story of 16-bit and 32-bit numbers from the number formats lesson. It feeds straight into Train a small LM, end to end, Fine-tuning and LoRA and Training on many GPUs.

Today

• Say, without looking anything up, why training a model takes about 16 bytes per parameter before counting anything that depends on the input, and name every one of those bytes.
• Explain in plain words why working backwards to find each weight’s correction needs the in-between results from the forward run, and estimate how big that pile gets from batch size, text length, model width and number of layers.
• Predict the peak memory of a real training step, then measure it on the Spark with PyTorch’s own memory counters and a recorded timeline of every allocation, and read the forward, backward and update phases off the picture.
• Cut that peak with three levers, one at a time: recomputing in-between results instead of storing them, storing the optimizer’s notes in 8 bits, and splitting a batch into smaller pieces whose corrections you add up; then say what each one costs you in speed.
• Answer “what’s the biggest model I can fully fine-tune on the Spark, and how much bigger if I only train small add-on pieces (LoRA)?” with arithmetic, not vibes.

Do How much memory does a model need? first: it covers the four places memory goes when serving, GB vs GiB, the difference between memory in use and memory PyTorch has set aside, and the Spark’s shared CPU-and-GPU memory. The number formats lesson helps for the precision part, but this lesson repeats what it needs. Plan on one evening of reading and arithmetic plus one session on the Spark.

The Problem

In How much memory does a model need? you learned to size inference: a 1B model in bf16 is about 2.5 GB of weights, plus a KV cache and some overhead. Serving Llama 3.2 1B fits on a laptop.

Now you want to fine-tune it. Same model, same 1.24 billion numbers. You write the obvious loop, batch of 8, sequence length 2,048, AdamW, on a 24 GB card, and it dies on the first step with something like:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.91 GiB.
GPU 0 has a total capacity of 23.55 GiB of which 1.12 GiB is free.

Nothing about the model changed. It is still 2.5 GB of weights. So where did the other twenty-plus gigabytes come from? You drop the batch to 1 and it runs, then dies on step 2, which is even stranger: what does step 2 need that step 1 did not?

Your cross-encoder at work has the same shape of problem. Whenever someone says “just train a bigger one”, the question that decides whether it happens on the Spark or on rented H100s is this lesson’s question.

The Idea

Four buckets, again, but different ones

At inference the kitchen had equipment bolted to the floor (weights), a tray per table (KV cache), a burst of counter space while plating (activations) and overhead. Training is a kitchen that is also rewriting the recipe after every dish. You still cook the dish (the forward pass). Then you go back through every step you took, asking “how should I have done this step differently?” (the backward pass). Then you change the recipe a little (the optimizer step). The extra memory is everything you need to keep so that “going back” and “changing a little” are possible.

GPU memory during training has these four buckets:

1. Weights. Same as before: parameters × bytes per parameter. Sometimes stored twice, in two precisions, for a reason we will get to.
2. Gradients. One number per parameter: “if you nudged this weight up, how much would the loss change?” The gradientgradientFor each parameter, how much the loss would change if that parameter grew a tiny bit. The backward pass computes it; the optimizer uses it to decide which way to move the parameter. has exactly the same shape as the weights, so it costs another full copy.
3. Optimizer states. What the optimizer remembers between steps. Plain SGD remembers nothing. Adam remembers two numbers per parameter. This bucket is usually the biggest.
4. Activations saved for backward. The intermediate values from the forward pass that the backward pass will need. At inference you threw them away the moment the next layer used them; in training you keep every layer’s worth until backward reaches it. This bucket grows with batch size and sequence length.

The first three are called model statesmodel statesThe memory that scales with the number of parameters during training: weights, gradients and optimizer states. Independent of batch size and sequence length., and they depend only on the model. The fourth depends on what you feed it.

Let’s take them one at a time, starting with the one that surprises people most.

The optimizer remembers two things about every weight

Plain gradient descent updates each weight as w = w − lr × grad. That has two problems in practice. Gradients from one batch are noisy, so the direction wobbles from step to step. And some parameters get huge gradients while others get tiny ones, so a single learning rate is too big for some and too small for others.

AdamAdamAn optimizer (Kingma and Ba, 2014) that keeps, for every parameter, a running average of its gradient (m) and a running average of its squared gradient (v), and moves each parameter by m divided by the square root of v. AdamW is Adam with weight decay applied separately; the memory is identical. fixes both by remembering two numbers per parameter:

• m, the first momentfirst momentA running average of the gradient, updated every step as m = 0.9 × m + 0.1 × grad. It smooths out the noise, like momentum on a ball rolling downhill.: a smoothed version of the gradient. “Which way has this weight wanted to move lately?”
• v, the second momentsecond momentA running average of the squared gradient, updated as v = 0.999 × v + 0.001 × grad². It measures how big this parameter’s gradients usually are, so Adam can scale its step.: a smoothed version of the gradient squared. “How big are this weight’s gradients, typically?”

The update divides one by the square root of the other, so every weight gets a step size scaled to its own history. That is why Adam works so well out of the box, and why it costs two extra full copies of the model. Both are kept in fp32 (4 bytes each) almost everywhere, because they are averages of many tiny numbers and lose accuracy fast in 16-bit. So Adam alone is 8 bytes per parameter, four times the bf16 weights it is updating.

I checked this on my Mac with a small GPT-2: before the first optimizer.step(), opt.state was empty. After it, every parameter had an exp_avg and an exp_avg_sq (that’s m and v in PyTorch’s naming), both fp32, both the parameter’s shape, totalling exactly 2.0 × the bytes of the fp32 weights. Note the after: PyTorch creates Adam’s states lazily, during the first step. That is the mystery from The Problem. Step 1 fit; step 2 did not, because step 2 is the first one that starts with m and v already in memory.

Why a second copy of the weights?

Modern training does the heavy math in bf16bf16Brain floating point: a 16-bit format with the same 8-bit exponent (range) as fp32 but only 7 stored mantissa bits (precision). Big range, coarse steps. Covered in the number formats lesson., because tensor corestensor coresDedicated matrix-multiply units on NVIDIA GPUs. They run 16-bit (and smaller) formats many times faster than fp32 math on the ordinary cores. run it several times faster than fp32 and it halves activation memory. But bf16 has a precision problem that bites at exactly one place: the update.

bf16 keeps about 8 significant bits, so neighbouring representable numbers near 1.0 are about 0.008 apart (2⁻⁷). A typical update is learning rate × something gradient-sized, maybe 0.0001 for a weight of size 0.05, a change of 0.2%. bf16 can only step in increments of about 0.4% to 0.8% of the value. Add 0.2% and the result rounds straight back to the original number. The update vanishes. Do that for every weight and training silently stops learning.

The fix, from the 2017 “Mixed Precision Training” paper (Micikevicius and others at NVIDIA and Baidu), is to keep a master weightsmaster weightsA full-precision (fp32) copy of the parameters that the optimizer updates. The 16-bit weights used in the forward and backward pass are re-copied from it after every step, so tiny updates accumulate instead of rounding away. copy in fp32. The optimizer adds the tiny update to the fp32 copy, where it survives. Then the bf16 copy is refreshed from it for the next forward pass. Many small updates add up in fp32 until they are big enough to show up in bf16.

The 16 bytes, and its variants

Put the classic recipe together, as the ZeRO paper counted it and as Megatron-style and DeepSpeed training lays it out:

So the rule of thumb: full training with Adam costs about 16 bytes per parameter, before activations. A 1B model is about 16 GB; Llama 3.2 1B at 1.24 billion parameters is 19.8 GB (18.4 GiB). An 8B model is 128 GB.

“About” matters, because the split depends on the setup, and you should know which one you are running:

• Plain fp32 training. Weights 4, gradients 4, m 4, v 4. Also 16. No master copy needed, because the weights are already fp32. Mixed precision does not save memory on model states at all. It saves activation memory and time.
• PyTorch AMP (torch.autocast with a normal fp32 model, the way most of your scripts work). The model’s parameters stay fp32 (4), so their gradients are fp32 (4), and Adam is fp32 (8): 16. Inside the autocast region PyTorch casts the weights to bf16 on the fly for each matmul and, by default, caches those casts until the region ends, which adds up to 2 more bytes per parameter at the peak. I’ll call it 16 + 2.
• Frameworks that also keep fp32 gradients. Megatron-LM, for example, can accumulate gradients in an fp32 buffer for accuracy across many micro-batches: 2 + 4 + 4 + 8 = 18.
• Pure bf16. Everything in bf16, no master copy: 2 + 2 + 2 + 2 = 8. Half the memory, and it works only with tricks to rescue the vanishing updates: stochastic roundingstochastic roundingRounding up or down at random, with probability proportional to how close the value is to each neighbour. A tiny update then survives on average instead of always rounding away., or a small Kahan-summation buffer that keeps the rounding error and adds it back next step. Expect worse results without them.
• Inference, for comparison: 2.

Why backward needs the forward’s leftovers

Now the fourth bucket. You built GPT-2 from scratch, so you have called loss.backward() many times. Here is what it needs.

The chain rulechain ruleThe calculus rule for nested functions: how much the output changes when you nudge the input equals the product of how much each step changes when you nudge the step before it., in plain words: the network is a long pipeline of small steps. To find how the loss depends on something early in the pipeline, you walk backwards from the loss, and at each step you ask “if my input moved a bit, how much would my output move?” Multiply those answers together and you have it. The backward passbackward passRunning the chain rule from the loss back to the first layer: each layer receives the gradient with respect to its output and computes the gradients with respect to its input and its own weights. is exactly that walk.

The catch is that each step’s “how much would my output move?” depends on what its input was. Three examples you have written:

• A linear layer, y = x @ W. The gradient for the weights is dW = xᵀ @ dy: you need x, the input this layer saw in the forward pass. (The gradient for its input, dx = dy @ Wᵀ, needs W, which you have anyway.)
• GELU (or ReLU, SiLU). Its slope depends on where on the curve you were. You need its input x.
• Softmax in attention. Its gradient is written in terms of its own output. You need the attention probabilities, unless the kernel is clever (more on that shortly).

So every operation, during forward, stashes whatever its own derivative will need. In Karpathy’s micrograd, that’s the _prev children and the values captured in each _backward closure. In PyTorch it is ctx.save_for_backward(...) inside every autograd function, and the stash is what autograd calls saved tensorssaved tensorsThe tensors an autograd operation keeps from the forward pass so it can compute its gradient in the backward pass. They are the “activations” in training memory.. Under torch.no_grad() nothing is stashed, which is why inference gets to throw activations away.

Look at the shape of that. The stash for layer 1 is written first and read last. At the moment the loss is computed, every layer’s stash is alive at once. That moment is the activation peak, and it’s why activation memory grows with depth.

How big the stash is

Count it for one transformer layer. Write b for batch size, s for sequence length, h for hidden size, a for attention heads, L for layers. One “sbh” is one hidden-sized vector per token in the batch, the basic unit of everything that flows through the model.

For a GPT-style block (4h-wide MLP with GELU, dropout on), Korthikanti and colleagues at NVIDIA counted, at 2 bytes per value:

• Attention: the input to the q/k/v projection, q and k for the score computation, v, the input to the output projection, a dropout mask: about 11 sbh bytes.
• MLP: its input, the 4h-wide pre-GELU and post-GELU values, a dropout mask: about 19 sbh.
• Two LayerNorms, each keeping its input: 4 sbh.
• Plus, in textbook attention, the score matrix: softmax output, its dropout mask and the dropped-out probabilities, 5 · a · s² · b bytes. Every head keeps an s × s grid for every sequence.

Total per layer, in bytes:

activations per layer  ≈  s · b · h · (34 + 5 · a · s / h)

and you multiply by L. The first term is linear in everything. The second grows with s², and it is the one that blows up at long sequences: at s = 2,048 with a = 32 heads and h = 2,048 it is 5 × 32 × 2,048 / 2,048 = 160, almost five times the first term.

Fused attention kernels (FlashAttention, and PyTorch’s scaled_dot_product_attention when it picks a fused backend) never store the s × s matrix. They keep only the output and one small number per row, and recompute the scores during backward, tile by tile. With them the s² term disappears, and per-layer activations are about 34 sbh. That’s the FlashAttention lesson seen from the training side: it saves more memory in training than in inference, because training would otherwise have to keep those matrices.

Two corrections to keep the formula honest:

• Llama-style blocks are bigger per token. No dropout, but the SwiGLU MLP is wide (8,192 for the 1B model, 4× hidden) and has three projections, and RoPE and RMSNorm save their own pieces. Counting the same way gets you roughly 40 to 50 sbh per layer. I use 46 for Llama 3.2 1B and 42 for Llama 3.1 8B in the figures; these are estimates.
• The logits are their own monster. At the end, the output head turns every token into one score per vocabulary word: b × s × V numbers. For GPT-2, V = 50,257; for Llama 3, V = 128,256. In bf16, then an fp32 copy for the loss, then the softmax’s saved output, then its gradient: call it 8 to 12 bytes per logit. For Llama 3.2 1B at b = 4, s = 2,048, that is 4 × 2,048 × 128,256 × 10 bytes ≈ 10 GB, as big as all sixteen layers together. Fused “chunked” cross-entropy kernels (Liger Kernel and others) exist precisely to avoid materialising this.

How close is the formula to reality? I counted every tensor autograd saved for Hugging Face’s GPT-2 small on my Mac’s CPU (fp32, so double the bytes; batch 2, sequence 1,024), using PyTorch’s saved_tensors_hooks (in torch.autograd.graph), which you’ll do yourself in the Build:

The formula predicts about 5.5 GB for the first row (including three fp32 copies of the logits). The measurement came in about 15% higher, probably mostly because HF’s GPT-2 computes its GELU as a chain of small Python-level ops, and each one saves its own 4h-wide intermediate. Treat the formula as a floor and measure your own model.

Watch one step

Now put all four buckets on a clock. Step through one training iteration: forward, loss, backward, the optimizer step and zero_grad. Then turn on checkpointing (explained next), and then look at the first step on its own.

Three things to notice:

1. The peak is at the turn, between forward and backward, when everything saved is alive at once. Not during the optimizer step, and not at the end.
2. Gradients and activations trade places. Backward frees layer L’s stash as it creates layer L’s gradients. Activations are usually much bigger than gradients, so memory falls through backward.
3. For GPT-2 small, the logits are the biggest single item. Its vocabulary (50,257) is 65 times its hidden size (768). For big models it’s the other way around, which is why advice about memory depends so much on model size.

The tools

Every trick for fitting training into memory removes or shrinks one bucket. Here they are, grouped by which bucket they attack.

Activations: gradient accumulation. If batch 32 doesn’t fit, run batch 8 four times, and only step the optimizer after the fourth. Gradients from each micro-batchmicro-batchOne slice of the batch that goes through forward and backward at once. Several micro-batches’ gradients are added together before one optimizer step. add up in the .grad buffers (that’s PyTorch’s default: backward() adds into .grad, it doesn’t overwrite). Divide the loss by 4 so the sum is an average, and the update is mathematically the same as one batch of 32 (up to small floating-point differences, and except for layers like BatchNorm that look across the batch). Activation memory is that of batch 8. Model states don’t change. The cost is time: four smaller passes keep the GPU less busy than one big one, so throughput drops a little, and a lot if the micro-batch gets so small that the GPU idles.

Activations: activation checkpointingactivation checkpointingAlso called gradient checkpointing. During forward, keep only some layers’ inputs and throw away everything else; during backward, re-run each layer’s forward from its saved input to rebuild the stash just before it’s needed.. Instead of storing every layer’s stash, store only each layer’s input (2 sbh in bf16), and during backward, re-run that one layer’s forward to rebuild its stash right before you need it. Memory drops from L full stashes to L small inputs plus one full stash. You pay with compute: every layer runs forward twice. A forward pass is roughly a third of a training step’s math (backward is about twice the forward: one matmul for the input’s gradient and one for the weight’s, against forward’s one), so the extra forward adds about 33% more compute. Measured slowdowns are usually 20 to 35%. In Hugging Face it’s model.gradient_checkpointing_enable(); in plain PyTorch it’s torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False). Chen and colleagues’ 2016 paper showed you can go further, checkpointing every √L layers for √L memory, but per-layer is what everyone uses.

Optimizer states: 8-bit optimizers. Store m and v in 8 bits instead of 32. Done naively that destroys them, because a single scale for a whole tensor can’t represent both the huge and tiny values. Tim Dettmers’ bitsandbytes stores them in block-wise quantizationblock-wise quantizationSplitting a tensor into small blocks (2,048 values in bitsandbytes’ optimizers) and quantizing each block with its own scale, so one outlier only hurts its own block., with a non-uniform 8-bit code that spends more of its 256 values near zero. Each value is turned back into fp32 just for the update and re-quantized after. m and v go from 8 bytes to 2 bytes per parameter, the paper reports no loss in quality on the tasks it tested, and it’s a one-line change: bnb.optim.AdamW8bit(model.parameters(), lr=...).

Optimizer states: Adafactor. Shazeer and Stern’s 2018 trick for v: for a weight matrix with n rows and m columns, don’t store n × m second moments. Store a running average for each row (n numbers) and each column (m numbers), and estimate any entry as row × column ÷ total. A 4,096 × 4,096 matrix needs 8,192 numbers instead of 16.8 million. By default it also drops m entirely. The optimizer states become a rounding error: in my check on a small GPT-2, torch.optim.Adafactor’s states were 0.2% of the weights. The trade is that it’s fussier to tune than Adam, and it’s most popular for T5-style models.

Weights and gradients: precision. Mixed precision makes the activations 2 bytes instead of 4, which often matters more than anything it does to the model states. Pure bf16 halves the model states too, at the price of the vanishing-update problem.

Everything: shard across GPUs (ZeRO, FSDP). With N GPUs doing plain data parallelismdata parallelismEvery GPU holds a full copy of the model, processes a different slice of the batch, and the GPUs average their gradients before each optimizer step., each GPU holds a full copy of all 16 bytes per parameter. That’s N copies of the same optimizer states, for no reason except convenience. ZeRO removes the redundancy in three stages: stage 1 splits the optimizer states (each GPU updates only its 1/N slice of the parameters), stage 2 also splits the gradients, and stage 3 splits the weights themselves, fetching each layer’s weights from the other GPUs just before it runs and dropping them right after. PyTorch’s version is FSDPFSDPFully Sharded Data Parallel: PyTorch’s implementation of ZeRO stage 3 style sharding. Each GPU stores 1/N of the parameters, gradients and optimizer states, and gathers full layers on the fly.. The distributed training lesson builds it; here is just the memory picture.

The ZeRO paper’s own example: a 7.5B model on 64 GPUs needs 120 GB per GPU of model states with plain data parallelism, 31.4 GB with stage 1, 16.6 GB with stage 2, and 1.9 GB with stage 3.

Everything: CPU offload. ZeRO-Offload and FSDP’s CPU offload keep optimizer states (and optionally weights) in the CPU’s RAM and run the Adam update on the CPU. On a normal server that’s a real win: an H100 has 80 GB, the box around it might have 1 to 2 TB of RAM. On the Spark it buys you nothing. As you saw in the virtual memory lesson, the Spark’s CPU and GPU share the same 128 GB. Moving a tensor “to the CPU” moves it to a different label on the same chips. The only offload that adds capacity on the Spark is to the SSD (DeepSpeed’s ZeRO-Infinity can offload to NVMe), and an SSD is hundreds of times slower than memory, as the latency lesson taught you to expect.

Put it together

Now play with every lever at once. The bar is the peak of one training step. The black lines are an H100 (sold as 80 GB; nvidia-smi reports about 79.6 GiB) and the Spark’s 128 GiB; the dashed line is what’s left on the Spark after the operating system, assumed here to take 10 GiB until you measure it.

Try these and watch which bucket moves:

• Llama 3.2 1B, BF16 + FP32 master, AdamW, batch 8, sequence 2,048: switch the optimizer between AdamW, 8-bit and Adafactor. Model states go 18.4 → 11.5 → 9.2 GiB.
• Same, now push the sequence to 16K. Then turn off fused attention. The s² term is why nobody trains long context without fused attention.
• GPT-2 124M at batch 64: the activations are more than 50 times the model states. For small models, the batch is the bill.
• Llama 3.1 8B with any setup: model states alone hit the Spark’s limit until you go to pure BF16 with 8-bit states or Adafactor.

What fits on the Spark

Do the arithmetic once, properly. Assume 118 GiB usable after the OS, and leave about 10 GiB for activations (small batches, checkpointing on). That leaves 108 GiB, about 116 GB, for model states.

These are upper bounds where “fits” means “runs at all”. Training speed is a separate question: the bandwidth lesson applies here too, because every optimizer step has to read all 16 bytes per parameter at least once (and write most of them back), and the Spark’s memory moves about 273 GB/s (a spec-sheet number). For a 7B model that’s at least 112 GB of traffic per optimizer step, 0.4 seconds or more before any of the forward and backward math.

LoRA changes the question. LoRALoRALow-Rank Adaptation (Hu et al., 2021). Freeze the model’s weights and train small extra matrices next to some of them: a d × r and an r × k matrix whose product is added to the frozen d × k weight. Only the small matrices get gradients and optimizer states. freezes the base model and trains small adapter matrices. The frozen weights need no gradients, no master copy and no optimizer states, so they cost just 2 bytes per parameter in bf16. The adapters get the full 16 bytes, but they’re tiny: rank 16 on every linear layer of Llama 3.1 8B is about 42 million parameters (0.5% of the model), 0.7 GB with Adam. So LoRA on 8B is about 16 GB of frozen weights plus 0.7 GB, plus activations.

And the activations don’t go away. Even though the base weights are frozen, the gradient still has to flow backward through every frozen layer to reach the adapters in earlier layers, and flowing backward through a layer needs its saved input. (You don’t need the dW half of the linear layer’s backward, so a bit less is saved, but most of the stash stays.) Checkpointing helps LoRA exactly as much as it helps full training.

So on the Spark: LoRA in bf16 reaches roughly 40-50B models (2 bytes per parameter into ~100 GiB). With QLoRAQLoRALoRA on top of a base model stored in 4-bit (Dettmers et al., 2023). The frozen weights are de-quantized to bf16 on the fly for each layer; adapters train in bf16. the base model is 4-bit, a little over 0.5 bytes per parameter with scales, so a 70B model’s frozen weights are roughly 35 to 40 GB and fit with plenty of room for activations. Full fine-tuning tops out around 7B; adapter fine-tuning reaches 70B. That is a 10× difference from the same 128 GB, and it is the main reason LoRA is the default way to fine-tune on one machine. The LoRA lesson builds it.

Build It

Everything runs on the Spark, in a folder like ~/workbench/trainmem/. The model is GPT-2 at four sizes, built from a config with random weights. Memory doesn’t care whether the weights are trained, so there is nothing to download and no gated model to request.

pip install transformers bitsandbytes
python -c "import torch, bitsandbytes; print(torch.__version__, torch.cuda.is_available(), bitsandbytes.__version__)"

# harness.py
import time, torch
from transformers import GPT2Config, GPT2LMHeadModel

SIZES = {  # n_embd, n_layer, n_head
    "124m": (768, 12, 12), "355m": (1024, 24, 16),
    "774m": (1280, 36, 20), "1558m": (1600, 48, 25),
}
GiB = 2**30

def build(size: str, ckpt: bool = False, n_positions: int = 1024):
    e, l, h = SIZES[size]
    cfg = GPT2Config(n_embd=e, n_layer=l, n_head=h, n_positions=n_positions,
                     resid_pdrop=0.0, embd_pdrop=0.0, attn_pdrop=0.0)
    model = GPT2LMHeadModel._from_config(cfg, attn_implementation="sdpa").cuda()
    if ckpt:
        model.gradient_checkpointing_enable()
    return model

def batch(b: int, s: int):
    x = torch.randint(0, 50257, (b, s), device="cuda")
    return x, x  # inputs, labels (the model shifts them)

def mem():
    return torch.cuda.memory_allocated() / GiB

def peak():
    return torch.cuda.max_memory_allocated() / GiB

# predict.py
def predict(params: int, h: int, L: int, V: int, b: int, s: int,
            recipe: str = "amp", optimizer: str = "adamw",
            ckpt: bool = False) -> dict[str, float]:
    """Return GiB for: weights, grads, optimizer, cast_cache, activations, logits, total.

    recipe: "amp" (fp32 params + bf16 autocast) is what your harness runs.
    optimizer: "adamw" | "adamw8bit" | "sgd"
    Activations: 34 * s * b * h bytes per layer (bf16, fused attention),
    or, with ckpt, 2 * s * b * h per layer plus one full layer.
    Logits: about 10 bytes per logit, b * s * V of them.
    """
    ...

# step.py
from harness import *

def step(model, opt, x, y) -> dict[str, float]:
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.synchronize(); t0 = time.perf_counter()
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = ...
    after_fwd = mem()
    ...  # backward, record; optimizer step, record; zero_grad(set_to_none=True), record
    torch.cuda.synchronize()
    return {"fwd": after_fwd, ..., "peak": peak(), "sec": time.perf_counter() - t0}

torch.cuda.memory._record_memory_history(max_entries=100_000)
for _ in range(3):
    step(model, opt, x, y)
torch.cuda.memory._dump_snapshot("snapshot.pickle")
torch.cuda.memory._record_memory_history(enabled=None)  # stop recording

def saved_bytes(model, x, y) -> int:
    seen: dict[int, int] = {}
    def pack(t):
        ...  # key by t.untyped_storage().data_ptr(), value storage nbytes; return t
    with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
        with torch.autocast("cuda", dtype=torch.bfloat16):
            model(x, labels=y)
    ...  # drop the parameters' own storages (saved weights are not activations)

for i, (x, y) in enumerate(micro_batches):
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(x, labels=y).loss / ACCUM
    loss.backward()            # adds into .grad
    if (i + 1) % ACCUM == 0:
        opt.step(); opt.zero_grad(set_to_none=True)

Measure

Fill this in for GPT-2 774M, s = 1,024, on the Spark. Use steps 3 to 5 for timing and peak (step 1 under-counts Adam; the first steps also pay one-off setup costs).

What the numbers should look like, and why. My predictions from the formulas, which you should check yourself:

Baseline. Model states 16 bytes × 774M = 11.5 GiB, plus 1.4 GiB of autocast cast cache, plus activations 34 × 1,024 × 8 × 1,280 × 36 bytes ≈ 12 GiB, plus logits 8 × 1,024 × 50,257 × 10 bytes ≈ 3.8 GiB: about 29 GiB. Expect the measurement to come in higher, perhaps 20 to 40%, because HF’s GPT-2 GELU is unfused (Task 4 shows you exactly how much). If your measured activations are far above that, check that SDPA really used a fused kernel; if dropout is on, it may not.

Step 1 vs step 2. The step 2 peak should be higher than step 1 by about 8 bytes × 774M = 5.8 GiB, Adam’s m and v, created in step 1’s optimizer.step() after the peak had already passed.

Checkpointing. Activations drop from ~12 GiB to about 1 GiB (36 layer inputs plus one rebuilt layer), so the peak falls by about 11 GiB to roughly 18 GiB, now dominated by model states and logits. Seconds per step should go up by 20 to 35%. If it went up much less, the step was probably not compute-bound to begin with.

8-bit AdamW. m and v go from 8 bytes to about 2 per parameter: about 4.3 GiB saved, no matter the batch. Step time roughly unchanged; the optimizer step is a small part of the total for a model this size.

Gradient accumulation. Activations and logits shrink by 4× (about 4 GiB instead of 16), model states unchanged. Seconds per optimizer step a little higher than baseline, because four small passes keep the GPU less busy than one big one. This is the only lever that doesn’t change the maths at all, only the schedule.

Max batch. Solve fixed + b × per_sequence ≤ 118 GiB. Baseline: about 14.5 GiB fixed and about 2 GiB per sequence, so roughly 50. Checkpointing cuts per-sequence cost to about 0.6 GiB, almost all of it logits, so roughly 170. 8-bit Adam only moves the fixed part, so it adds just two or three sequences. Your measured numbers will be lower; the ratios between configurations are what to check.

Then check the scaling, the interview question. Repeat the baseline at s = 512, 1,024 and 2,048 (with n_positions=2048). With SDPA, activations should double each time. Then switch to eager attention and repeat: they should grow faster than double, because of the s² score matrices.

Break It

1. Turn off fused attention at long sequences. Use attn_implementation="eager" at s = 4,096 with GPT-2 124M, batch 4 (build it with n_positions=4096). Predict the score-matrix bytes first: 5 × 12 heads × 4,096² × 4 × 12 layers. That is about 48 GB for score matrices alone, for the smallest GPT-2. Watch it fail where SDPA ran comfortably. Lesson: in training the s² term isn’t a spike you pass through, it’s a stash you keep until backward.

2. Forget zero_grad. Comment it out and run five steps. Memory doesn’t keep growing (gradients accumulate into the same .grad buffers), but the loss goes haywire because every step now uses the sum of all previous gradients. Then put it back as zero_grad(set_to_none=False), which fills the buffers with zeros instead of freeing them. Training is correct again, but compare the peak with the default: it’s higher by one full gradient copy (4 bytes × 774M ≈ 2.9 GiB under AMP), because the gradients now sit in memory through the next forward pass, right when activations peak. Find the difference in your snapshot. That’s why PyTorch made set_to_none=True the default.

3. Keep the graph alive by accident. Add losses.append(loss) (not loss.item()) inside the loop and run fifty steps. Memory climbs every step: each stored loss tensor holds a reference to its whole autograd graph, and with it every saved activation of that step. This is the most common real-world “memory leak” in training loops. Fix it with loss.item() or loss.detach(), and see the climb stop.

4. Train in pure bf16 without a master copy. Convert the model with model.to(torch.bfloat16), use plain AdamW, and set a small learning rate like 1e-5. Train on a tiny fixed batch you can overfit. Compare the loss curve to the AMP run. With a small learning rate, many updates round to nothing in bf16 and the loss stalls or falls far more slowly. Memory for model states halves; this is the price.

Explain It Back

Answer out loud, in plain words, as if in an interview. Then check.

Go Deeper

• Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2019). Section 3 is the clearest accounting of training memory ever written: the 16 bytes, the activation estimate, fragmentation, and the three stages.
• Paulius Micikevicius et al., Mixed Precision Training (2017). Why the fp32 master copy exists, with the plot of updates vanishing below fp16’s precision. Short and very readable.
• Vijay Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models (2022). Where the 34 sbh + 5as²b formula comes from, term by term, and “selective” checkpointing that recomputes only the cheap-to-recompute attention parts.
• Tianqi Chen, Bing Xu, Chiyuan Zhang, Carlos Guestrin, Training Deep Nets with Sublinear Memory Cost (2016). The original gradient checkpointing paper, with the √n trick.
• Tim Dettmers et al., 8-bit Optimizers via Block-wise Quantization (2021). How bitsandbytes keeps Adam’s states in 8 bits without hurting training.
• Noam Shazeer and Mitchell Stern, Adafactor: Adaptive Learning Rates with Sublinear Memory Cost (2018). The row-and-column trick for the second moment.
• PyTorch docs, Understanding CUDA Memory Usage. The official guide to _record_memory_history, snapshots and the memory_viz tool you use in Task 3.
• Yanli Zhao et al., PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023). How ZeRO-3 became part of PyTorch, and the engineering it took. Good preparation for the distributed lesson.

Connections

This lesson extends How much memory does a model need?: same bytes-per-parameter thinking, same allocator, same unified-memory caveats, with three new buckets that inference gets to skip. The s² activation term is the training half of FlashAttention, and the fp32 master copy is number formats applied to a real failure.

It unlocks the next three lessons in this track. In Train a small LM, end to end you’ll pick batch size, sequence length and checkpointing for a ~100M model on the Spark using exactly this calculator. LoRA builds the adapter trick that turns “7B, at most” into “70B, comfortably”. And distributed training builds the ZeRO picture from Figure 005: when one machine’s 128 GB really isn’t enough, you split the 16 bytes across many.