Training memory, explained

Where DNN memory goes

Training a network puts four kinds of tensors on the GPU: weights, activations, gradients and optimizer state. They are created and freed at different moments, and each one grows with something different. This page shows when each exists and lets you size them for your own model.

One parameter, trained with mixed-precision Adam, costs 16 bytes.

Weight in bf16 2 B
Master weight in fp32 4 B
Adam m and v in fp32 8 B
Gradient in bf16 2 B

Each cell is one byte. This is the cost before a single activation is stored. A 7 billion parameter model needs about 112 GB for these four items alone.

Four kinds of memory

Everything the framework allocates for training falls into one of these buckets, plus a few temporary buffers. The colors below are used for the same thing everywhere on the page.

Weights

What: The learned numbers: every matrix W and bias b in the network.
Lives: The whole run. Loaded once, nudged every step, saved in checkpoints.
Size: Parameter count × bytes per parameter. The same at batch size 1 or 1,000.

Activations

What: The output of each layer for the current batch, such as a₁ = f(W₁a₀).
Lives: Created in forward, kept until backward uses them, then freed.
Size: Grows with batch size × sequence length × number of layers.

Gradients

What: ∂Loss/∂W for every weight. One number per parameter, shaped like the weights.
Lives: Created in backward, read by the optimizer step, cleared before the next batch.
Size: Equal to the weights, in the same data type.

Optimizer state

What: Extra numbers the optimizer keeps per weight. Adam keeps a running mean m and variance v.
Lives: The whole run, from the first update onward.
Size: Plain SGD 0, momentum 1×, Adam 2× the parameter count, usually in fp32.

The word "activation" is easy to misread. An activation function such as ReLU is an operation and holds no memory. The activations are the tensors it produces, and those are what backward needs later. Other buffers also exist: activation gradients that flow backward for a moment, kernel workspaces, the input batch, and at inference time the KV cache. They sit in the calculator under temporary buffers and KV cache.

Why backward needs the forward pass

Layer i computes z_i = W_i · a_i−1 and then a_i = f(z_i). During backward, the error arriving at that layer is δ_i = ∂L/∂z_i. Two products come out of it.

Gradient for the weights

∂L/∂W_i = δ_i · a_i−1^T

Needs the activation a_i−1 that layer i saw on the way forward. That is why forward cannot throw its outputs away.

Error passed to the previous layer

∂L/∂a_i−1 = W_i^T · δ_i

Needs the weights W_i, so they stay resident throughout backward.

The first product has one entry per weight, so gradients match the weights in size. The saved activations have one entry per layer output per sample, so they grow when you add samples or tokens. That difference explains most memory behavior you will see below.

One training step, tensor by tensor

A three-layer network, nine moments. F is forward through a layer, B is backward through a layer. Drag the slider or click a column. A colored cell means the tensor is in memory at the end of that moment.

Weights
Optimizer state
Saved activations
Temporary (δ)
Gradients (∇W)

Moment in the step

Allocated

Freed

Size it for your model

Choose a transformer and a batch. The chart below uses the same nine moments as above, with the network split into thirds. Results are estimates in GiB (1 GiB = 1,024³ bytes) that follow the formulas in the papers listed at the end.

Run

Training Inference

Model Preset

Layers Hidden size Attention heads Vocabulary

Tied input and output embeddings

Batch

Sequence length Batch size per GPU

Numbers Precision Optimizer

Memory savers FlashAttentionNever stores the s × s attention scores. Activation checkpointingKeeps only each layer's input and recomputes the rest in backward.

Compare with GPU memory

Show the numbers

Memory ledger

The largest size each kind reaches, and the rule that produced it for your settings.

Levers that move each bucket

Each technique attacks one bucket and pays for it somewhere else.

Lever	What shrinks	What it costs
Smaller batch per GPU	Saved activations, in proportion to batch size	Lower GPU utilization, or more gradient accumulation steps
Gradient accumulation	Nothing directly. It gives a large effective batch while holding activations at the small batch size	More forward and backward passes per update
Activation checkpointing	Saved activations drop to one input per layer plus one layer being recomputed	About one extra forward pass of compute
FlashAttention	The s × s attention scores, which otherwise grow with the square of the sequence length	None in practice, it is a fused kernel
Sharding (ZeRO-3, FSDP)	Weights, gradients and optimizer state, divided across GPUs	Network traffic between GPUs
8-bit or bf16 optimizer state	Optimizer state, by 2× to 4×	Some risk to training stability
LoRA or frozen layers	Gradients and optimizer state exist only for the trainable parameters	Less capacity to change than full fine-tuning
CPU or disk offload	Optimizer state and weights move to host memory	Slow transfers over PCIe

To check an estimate against a real run, reset the peak counter, run one full training step, and read it back.

torch.cuda.reset_peak_memory_stats()
# one full step: forward, backward, optimizer.step(), optimizer.zero_grad()
print(torch.cuda.max_memory_allocated() / 2**30, "GiB")

The CUDA context (a few hundred MiB) and allocator fragmentation sit on top of that number in nvidia-smi.