Where memory goes: training vs inference

This is the chapter the whole book's constraint hangs on. Sixteen gigabytes is a hard wall, and the difference between a lab that runs and a lab that dies with CUDA out of memory is almost always a memory-accounting mistake, not a compute one. So I am going to account for every byte, in both regimes. Inference memory is short: weights plus the KV cache, and that is essentially it. Training memory is the interesting one, because reverse-mode autograd (from chapter 1) forces you to store activations, and the AdamW optimizer keeps two extra full-size copies of every parameter, so training a model can cost eight to sixteen times what running it costs. I will derive the backprop-through-a-linear-layer rule that makes the activation-storage requirement concrete, write down the memory-accounting formulas, and then do the byte arithmetic for a 4B model in both modes on the RTX 5080, which is exactly the calculation that tells you what you can and cannot fine-tune on this card.

Theory

Inference memory: weights plus KV cache

At inference, with torch.no_grad() on, autograd records nothing, so there are no stored activations to speak of beyond the single layer being computed. The memory is two things. First, the weights, which is just parameter count times bytes-per-element: a model with $P$ parameters in a dtype of $b$ bytes needs $P\cdot b$ bytes resident. Second, the KV cache, the stored keys and values for every past token so the model does not recompute them each step, whose size I derived in the transformer-block chapter (equation 4.9):

$$M_{\text{infer}}=\underset{\text{weights}}{\underbrace{P\cdot b}}+\underset{\text{KV cache}}{\underbrace{2\cdot B\cdot L\cdot N\cdot h_{kv}\cdot d_{\text{head}}\cdot b}}+(\text{small: current activations, CUDA context}).\text{\hspace{1em}(6.1)}$$

The activations at inference are negligible because you only ever hold one layer's worth at a time (the residual stream is a single $B\times L\times d$ tensor passing through), and the CUDA context and kernels are a fixed few hundred megabytes. So inference is dominated by weights, with the KV cache growing linearly in sequence length and becoming the swing factor at long context. This is why the whole of Part II is about serving: on a fixed card, weights are fixed but the KV cache is where you win or lose, and GQA (equation 4.10) was the first lever for it.

Training memory: why it explodes

Training adds three new consumers, and together they are why training costs multiples of inference. The first is gradients: every trainable parameter needs a same-shaped gradient tensor, another $P$ elements. The second is optimizer state, and for AdamW this is the killer. The third is stored activations, forced by reverse-mode autograd. Take them in order, starting with the one that requires a derivation.

Now the accounting. Consider a full fine-tune with mixed-precision AdamW, the standard recipe. Per parameter, the framework holds: an FP32 master copy of the weight (4 bytes, kept in full precision so tiny updates are not lost to BF16's coarse mantissa, per chapter 1's equation 1.8), a BF16 working copy used in the forward and backward passes (2 bytes), a gradient (2 bytes in BF16, sometimes 4), and AdamW's two optimizer moments, the first moment $m$ and the second moment $v$, each an FP32 same-shaped tensor (4 bytes each). That is the famous accounting:

$$\underset{\text{fp32 master}}{\underbrace{4}}+\underset{\text{bf16 weight}}{\underbrace{2}}+\underset{\text{grad}}{\underbrace{2}}+\underset{\text{Adam }m}{\underbrace{4}}+\underset{\text{Adam }v}{\underbrace{4}}=16\text{ bytes per parameter}.\text{\hspace{1em}(6.4)}$$

AdamW is "2 extra copies" in the sense that $m$ and $v$ are each a full parameter-sized FP32 tensor on top of the weights and gradient; those two moments alone are 8 of the 16 bytes. So the fixed, activation-independent cost of full-fine-tuning a $P$-parameter model is $16P$ bytes, versus $2P$ bytes to merely run it in BF16. That factor of eight, before activations, is most of the "8-16x more" this chapter opened with. On top of it sit the stored activations from equation (6.3), whose size scales with batch size, sequence length, and depth:

$$M_{\text{train}}=\underset{\text{weights, grads, Adam }m,v}{\underbrace{16P}}+\underset{\text{stored activations}}{\underbrace{A\cdot B\cdot L\cdot N}}+(\text{CUDA context}),\text{\hspace{1em}(6.5)}$$

where $A$ is the per-token, per-layer activation footprint (a handful of $d$-sized and $d_{\text{ff}}$-sized tensors that must survive to backward). The activation term is the one you can actually control at run time, which is where the next idea comes in.

Gradient checkpointing

Equation (6.5)'s activation term grows with depth $N$ and can dominate for long sequences or large batches, and it is the one term you can trade against compute. Gradient checkpointing (activation checkpointing) is the deal: during the forward pass, do not store the intermediate activations for every layer; store only a sparse set of "checkpoints" (say, one per layer boundary or every $\sqrt{N}$ layers). Then during the backward pass, when equation (6.3) needs the input $X$ to a layer that was not stored, recompute it by re-running the forward pass from the nearest checkpoint. You pay one extra forward pass worth of compute (roughly a 30 percent wall-clock hit) in exchange for cutting the stored-activation memory from $O(N)$ down to $O(\sqrt{N})$ or less. It is the purest expression of the memory-versus-compute tradeoff that reverse-mode autograd sets up: autograd wants to store everything (chapter 1), checkpointing declines to, and recomputes on demand. On a 16GB card it is frequently the difference between a batch size of 1 and a usable batch size, and every training lab in this book turns it on.

The bigger structural escape, which I only flag here and derive in the post-training part, is to not update most of the parameters at all. LoRA and its quantized cousin QLoRA freeze the base weights (so they need no gradient, no master copy, and no Adam moments, collapsing the $16P$ term to $2P$ for the frozen base plus $16\times (\text{tiny})$ for the small adapter) and keep the frozen base in 4-bit (halving even that $2P$). That is how a model whose full fine-tune needs many times the card's capacity still trains on it, and equation (6.4) is exactly the cost structure LoRA is dodging.

Tooling

The tool is PyTorch's CUDA caching allocator and its introspection API, plus the torch.utils.checkpoint module for gradient checkpointing. The allocator does not hand each tensor its own cudaMalloc; it grabs large blocks from the driver and sub-allocates, which is why freed tensors do not immediately return memory to the OS (they return to PyTorch's pool) and why nvidia-smi shows a higher, stickier number than your tensors' live bytes. The functions that tell the truth about your program's own usage are torch.cuda.memory_allocated() (bytes currently held by live tensors), torch.cuda.max_memory_allocated() (the peak, which is what actually has to fit), and torch.cuda.memory_summary() (a full human-readable table broken down by allocation category). The distinction between allocated and reserved is the one to internalize: reserved is what PyTorch has taken from the driver, allocated is what your tensors are using, and the gap is pool headroom, not a leak. For checkpointing, torch.utils.checkpoint.checkpoint wraps a submodule so its activations are dropped and recomputed, and most model libraries expose a gradient_checkpointing_enable() one-liner that flips it on for every block. The whole discipline of training on 16GB is: watch max_memory_allocated, keep it under the card's capacity, and reach for checkpointing, quantization, and LoRA (equations 6.4–6.5) when it climbs too high.

Lab

The goal is to measure the memory accounting of equations (6.1)–(6.5) on a real small model: load it for inference and read the footprint, then run a training step (forward, loss, backward, optimizer step) and watch the peak climb through the predicted multiple as gradients, optimizer state, and stored activations come online. The artifact is a JSON report of measured peaks in each phase, checked against the byte arithmetic.

uv init labs/memory
cd labs/memory
uv add torch transformers

"""Measure inference vs training memory and reconcile with the formulas.

Artifact: artifacts/memory_report.json with per-phase peak allocations.
Run on a CUDA machine; on CPU it reports parameter math only.
"""
import json
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "Qwen/Qwen3-0.6B"
OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

def gb(nbytes: int) -> float:
    return round(nbytes / 1e9, 3)

def main() -> None:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tok = AutoTokenizer.from_pretrained(MODEL)
    P = None
    report = {"model": MODEL, "device": device}

    # --- Inference footprint (eq 6.1). ---
    model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
    P = sum(p.numel() for p in model.parameters())
    report["params"] = P
    report["predicted_weights_bf16_GB"] = gb(P * 2)
    report["predicted_full_train_16B_GB"] = gb(P * 16)

    if device == "cuda":
        model = model.to("cuda")
        torch.cuda.reset_peak_memory_stats()
        ids = tok("Explain entropy in one sentence.", return_tensors="pt").input_ids.to("cuda")
        with torch.no_grad():
            model(ids)
        report["measured_inference_peak_GB"] = gb(torch.cuda.max_memory_allocated())

        # --- One training step (eq 6.4/6.5). ---
        model.train()
        opt = torch.optim.AdamW(model.parameters(), lr=1e-5)  # allocates m, v
        torch.cuda.reset_peak_memory_stats()
        out = model(ids, labels=ids)
        out.loss.backward()          # allocates grads + stored activations (eq 6.3)
        opt.step()                   # materializes Adam m, v (eq 6.4)
        report["measured_train_step_peak_GB"] = gb(torch.cuda.max_memory_allocated())
        report["train_to_infer_ratio"] = round(
            report["measured_train_step_peak_GB"]
            / report["measured_inference_peak_GB"], 2)
        print(torch.cuda.memory_summary(abbreviated=True))
    else:
        report["note"] = "No CUDA; reporting parameter arithmetic only."

    (OUT / "memory_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))
    print(f"Artifact: {(OUT / 'memory_report.json').resolve()}")

if __name__ == "__main__":
    main()

uv run python measure_memory.py

What you should see. The script writes artifacts/memory_report.json with the parameter count, the predicted BF16 weight footprint ($2P$) and full-train footprint ($16P$), and the measured peaks. On the GPU, the inference peak lands near the predicted weight bytes plus a little context and KV, while the training-step peak is a clear multiple of it (the train_to_infer_ratio field), because backward brought gradients and stored activations (equation 6.3) online and opt.step() materialized Adam's moments (equation 6.4). Watching that ratio jump from roughly 1 to several in a single script is the entire chapter made empirical: nothing about the model changed, only whether you asked autograd to remember the forward pass and whether an optimizer allocated its state. Record the exact peaks with date and driver (measured on the baseline machine, record value, date, driver); the memory_summary() table printed alongside shows the allocated-versus-reserved split so you can see the caching allocator's headroom directly.