Memory accounting before memory tricks

An out-of-memory error does not identify which allocation dominates. Model parameters, gradients, optimizer state, activations, temporary buffers, and allocator behavior all contribute. A rough accounting model helps decide what to measure next. It also prevents choosing a technique that targets the wrong term.

Begin with parameter count and bytes per stored value. Then state the optimizer and precision assumptions. A simple full-precision Adam-style estimate includes parameters, gradients, and two moment arrays, often four parameter-sized float32 arrays before other overhead. Mixed-precision implementations may retain master weights or use different state layouts, so inspect the actual stack rather than treating one formula as universal.

Activations often scale with batch size, sequence length, and model structure. They can dominate even when parameter storage is fixed. Activation checkpointing trades recomputation for lower saved activation memory. Gradient accumulation reduces per-microbatch activation needs but does not inherently reduce parameter or optimizer storage. Sharding targets replicated states across devices. Each technique changes a different term.

Measure peak memory over a representative step, including backward and optimizer update. A snapshot after forward may miss the peak. Distinguish allocated tensor memory from reserved allocator memory. Avoid claiming that every unused reserved byte is a leak. Repeated growth across steps and retained computation graphs need a separate investigation.

A fictional model has 100 million parameters. Under a deliberately simple float32 Adam accounting model, each parameter-sized array uses 400 MB in decimal units. Parameters, gradients, and two moments total 1.6 GB. Suppose measured activations peak at 6 GB and temporary buffers at 1 GB. The first-order total is 8.6 GB before additional framework overhead.

Halving the microbatch may reduce the activation term approximately, but leaves the 1.6 GB state term mostly unchanged. If activations halve to 3 GB, the estimate becomes 5.6 GB. This is a scenario estimate, not a guarantee. Variable sequence lengths and operator workspaces can break simple linear scaling.

A model has 50 million parameters under the same four-array float32 assumption. Activations use 4 GB. Calculate state memory and explain whether activation checkpointing or optimizer sharding more directly addresses the larger term.

State memory is 800 MB decimal. Activations are larger, so activation checkpointing targets the dominant stated term, subject to runtime tradeoffs. Optimizer sharding could still help in another configuration but addresses a smaller part here. Award one point each for the arithmetic, dominant-term identification, technique-to-term mapping, and measuring the resulting peak. Do not claim that a technique is free; report its compute or communication cost.

Decimal gigabytes use one billion bytes; binary gibibytes use 2^30 bytes. Keep the unit explicit when comparing an estimate with a device capacity or allocator report. The initial 1.6 billion bytes of state are approximately 1.49 GiB. A mismatch in units can matter near the device limit, but it does not explain every large allocation gap.

This additive table is a first-order scenario, not automatically a valid measured peak. Two allocations with separate maxima may not coexist. Conversely, optimizer temporaries may appear during update after forward activations were freed. For a rigorous peak ledger, inspect the timeline and count live allocations at the same instant. The table tells you which terms to investigate.

A training loop appends the loss tensor to a history list without detaching it. Depending on how the tensor and computation graph are retained, references can keep graph-related objects alive. A more appropriate logging path stores a detached scalar at the needed frequency and checks memory over repeated steps. Do not diagnose every upward reserved-memory curve as a retained graph: allocators can keep reusable blocks after tensors are freed.

1Step                 1      2      3      42Allocated after step 3.0    3.0    3.0    3.0 GB3Reserved after step  4.0    5.0    5.0    5.0 GB4Interpretation: stable live tensors; reservation growth alone is not a leak56Alternative trace7Allocated after step 3.0    3.6    4.2    4.8 GB8Interpretation: persistent live growth needs an ownership/reference audit

The alternative trace does not prove that the history list is responsible. Compare runs with logging disabled or with detached logging, inspect retained tensors, and localize the first step that keeps unexpected state. If memory rises only when batch sequence length grows, the workload itself may explain the pattern.

A model has a six-gigabyte replicated state footprint, a two-gigabyte activation peak, and one gigabyte of temporary buffers. You can approximately halve activations through checkpointing, or distribute four gigabytes of optimizer state evenly across four workers with extra communication. Under a deliberately simple no-overlap accounting model, checkpointing saves one gigabyte. Ideal optimizer-state distribution leaves one of the four optimizer gigabytes per worker, saving three gigabytes per worker and reducing the stated total from nine to six.

Neither estimate guarantees the real peak: communication buffers, gathered states, recomputation workspaces, and framework layout can change it. Award one point for each saving, one for the targeted term, and two for the limitations and required peak remeasurement. The option with the greatest memory reduction may not give the best throughput or engineering cost.

“Gradient accumulation reduces all training memory by the accumulation factor” fails because parameter and optimizer state remain, and some implementations retain graphs incorrectly. “Activation checkpointing is a free memory reduction” fails because it repeats computation and can interact with randomness and implementation choices. Check the framework's contract and compare both correctness and runtime.

Keep a maximum-input fixture for variable-length workloads. An average sequence may fit while the longest allowed batch fails. A practical release decision needs either a supported shape envelope or a defined split/reject route for larger inputs. Report the envelope with the peak. A benchmark that quietly shortens sequences has changed the task and cannot support an equal-work memory claim.

Original practice: Will gradient accumulation fix every memory failure? A strong answer distinguishes activation memory from parameters and optimizer state. Follow up with a failure during optimizer step. A weak answer reduces batch size without locating the peak.