Training memory accounting

Why it matters

Whether a run fits in memory is decided before the first step, by arithmetic you can do on paper. Interviewers ask "will this model train on these GPUs, and if not, what would you change?" to check whether you estimate before running and whether you know which lever moves which term.

Core ideas

Model states scale with the parameter count Ψ. Count bytes per parameter:

• fp32 training with Adam: 4 (weights) + 4 (gradients) + 8 (two moments) = 16 bytes.
• 16-bit compute with fp32 master weights and Adam, the setup analyzed in the ZeRO paper: 2 (16-bit weights) + 2 (16-bit gradients) + 4 (fp32 master copy) + 8 (two fp32 moments) = 16 bytes.
• SGD with momentum in fp32: 4 + 4 + 4 = 12 bytes; plain SGD keeps no optimizer state.

So mixed precision with master weights does not shrink model states; its savings are in activations and speed. As simple arithmetic, a 1-billion-parameter model at 16 bytes per parameter needs about 16 GB for model states before any activation is stored. Frameworks differ (some keep fp32 gradients, some train purely in bf16 without a master copy), so check what your setup actually stores.

Activations scale with the micro-batch. The backward pass needs values from the forward pass: inputs to linear layers, attention probabilities, normalization statistics, dropout masks. For a transformer, activation memory grows roughly linearly with micro-batch size b, sequence length s, hidden size h and layer count L, plus a per-layer term proportional to b × heads × s² when attention scores are materialized. That is why runs often fit at short context and fail at long context with the same model.

Everything else. Short-lived temporaries, library workspaces, gradient buckets for all-reduce, the output logits (b × s × vocabulary, plus the softmax for the loss, often a spike with large vocabularies), and fragmentation in the caching allocator. PyTorch reports allocated memory (live tensors) separately from reserved memory (held by the allocator); nvidia-smi shows reserved memory plus the CUDA context, so it overstates live tensors.

Peak, not average. Memory peaks at a specific moment, often early in the backward pass when saved activations and new gradients coexist, or in the first optimizer step when optimizer state is created lazily. Measure the peak.

model_states = params * bytes_per_param(precision, optimizer)
activations  ≈ layers * per_layer_activations(micro_batch, seq_len, hidden, heads)
peak         ≈ model_states + activations + largest_temporary (e.g. logits)
               + communication_buffers + allocator_slack

Then run one step, read the peak allocated memory and a memory snapshot, and explain the gap.

Levers, roughly from cheapest to most invasive: a smaller micro-batch with gradient accumulation (activation memory of one micro-batch, same effective batch); fused attention kernels that never store the s×s matrix; activation checkpointing (store less, recompute in backward); sharding model states across data-parallel ranks (ZeRO stage 1 shards optimizer state, stage 2 adds gradients, stage 3 or FSDP adds parameters); tensor or pipeline parallelism; and CPU or NVMe offload, limited by transfer bandwidth. Lower-precision optimizer states also exist but change numerics.

Common misconceptions

• "Mixed precision halves training memory." With fp32 master weights and Adam, model states stay near 16 bytes per parameter; the savings are in activations.
• "Memory grows with the global batch size." It grows with the micro-batch; gradient accumulation keeps the global batch while bounding activation memory.
• "nvidia-smi tells me what my tensors use." It reports reserved memory and the CUDA context; use allocated and peak-allocated counters and memory snapshots.
• "Activation checkpointing doubles the step time." Recomputing the whole forward pass costs about one extra forward, roughly a third more compute when backward costs about twice the forward; selective recomputation costs less.
• "Inference needs training memory minus gradients." Inference drops optimizer state and most activations but adds the KV cache, which can dominate at long contexts and high concurrency.

How to approach it

Write the estimate before running anything, then measure and explain every difference. In interviews, lay out the per-parameter accounting, then activation scaling, then the levers in order of cost. Senior candidates are expected to add communication buffers and fragmentation, and to discuss how sharding, recomputation and throughput interact.