Memory, Throughput, and Compute-Efficient Training

Training memory has several owners, and which one dominates changes with model, sequence length, batch, and optimizer. For mixed-precision Adam the split has been written down byte by byte. It comes to 16 bytes per parameter: 2Ψ for the fp16 weights, 2Ψ for the fp16 gradients, and 12Ψ for the fp32 parameter copy, momentum and variance. Twelve of those sixteen bytes are optimizer state.

The ZeRO paper says what that costs on a real model: “For a model such as GPT-2 with 1.5 Billion parameters, this leads to a memory requirement of at least 24 GB, which is significantly higher than the meager 3 GB of memory required to hold the fp16 parameters alone.”

And that is only the state. The same paper puts the activations for that same 1.5B-parameter GPT-2, at sequence length 1K and batch size 32, at about 60 GB more. A fifth owner, larger than the four before it.

Hugging Face's Transformers documentation reaches the same anatomy independently, and per parameter: 6 bytes of mixed-precision weights (fp16 plus the fp32 master copy), 8 bytes of Adam momentum and variance, 4 bytes of fp32 gradients, plus activations that “vary in size with batch size, sequence length, model depth, and hidden size”. Two independent accounts, one conclusion. The weights are rarely the thing that fills the card.

Each of these four saves one resource by spending another. Choose the trade that fits the limiting resource, and where someone has measured the trade, use their measurement rather than an adjective.

Take the fourth column. Padding is not a rounding error. Graphcore's packing paper opens with the size of the waste: “We show in this paper that the variation in sequence lengths in common NLP datasets is such that up to 50% of all tokens can be padding.” For GLUE-CoLA at sequence length 128 the figure reaches 89%. Removing that waste, without letting packed examples see one another, gave a 2x speedup for BERT phase-2 pre-training on Wikipedia at sequence length 512.

The “Verify: attention boundaries” line in the same column is not a precaution invented for a course. Earlier packing implementations “did not consider example boundaries when using Flash Attention 2, resulting in undesired cross-example attention that reduce quality and convergence”. That was reported on both the Hugging Face and IBM Research blogs. Boundary-aware packing — passing position_ids and calling flash_attn_varlen_func — bought up to 2x throughput and a 20% peak-memory reduction on FLAN, and 1.4x on OrcaMath, measured in tokens per second across llama2-7B, mistral-7B and granite-8B-code on 8 A100-80 GPUs.

The saving and the failure mode are the same size. The same 2x sits on either side of a correct mask.

A local metric may improve while the total system becomes slower or riskier. MLPerf, the field's training benchmark, says so in the paper that defines it: “Although many optimizations immediately improve traditional performance metrics such as throughput, some can decrease the final model quality, an effect that is only observable by running an entire training session.” The abstract states the reversal plainly: “optimizations that improve training throughput can increase the time to solution”.

The cost did not vanish. It moved somewhere the throughput counter does not look.

A GPU can wait for data decoding, host transfers, distributed synchronization, checkpoint writes, or Python control flow. That wait has been measured at fleet scale rather than guessed at. Google analysed millions of tf.data jobs running in its own datacenters and concluded: “On average, ML training jobs spend 30% of their total compute time in the input pipeline.” The tail is worse than the mean. The same analysis reports that “20% of jobs spend more than a third of their compute time in the input pipeline”, and that 13% of jobs read at least 1 TB of input data.

Microsoft Research and UT Austin found the same thing independently, across nine models, three tasks and four datasets on production clusters: “We find that in many cases, DNN training time is dominated by data stall time: time spent waiting for data to be fetched and preprocessed”. Their CoorDL loader cut training time “by as much as 5x on a single server” against DALI. Two fleets, two teams, one bottleneck — and it is not the kernel.

The mirror case is real too, and it is the one profiling misses. The arithmetic units can sit near peak while end-to-end time stays poor, because the step is waiting somewhere the FLOP counter does not look.

So profile representative steady-state and transition phases. Include evaluation, accumulation, compilation, and data-loader warmup rather than timing one ideal batch.

Architecture and dataset choices can dominate kernel-level optimizations. A smaller model trained on cleaner, better targeted evidence may reach the product threshold sooner.

That budget question was tested directly in 2022. Hoffmann and twenty-one colleagues trained “over 400 language models ranging from 70 million to over 16 billion parameters”. The rule they extracted from those runs is a short one: for compute-optimal training “the model size and the number of training tokens should be scaled equally”.

Chinchilla holds 70B parameters. It outperformed Gopher, which holds 280B. The smaller model was the better use of an identical compute budget.

Use representative shapes and full training state. Each step below is a measurement, not an impression. The last one is the only step whose number a benchmark would accept.

A system may report high throughput while processing padding, duplicated examples, masked positions, or low-value easy data. Raw throughput does not measure useful evidence, and the industry's own training benchmark does not score it: “MLPerf’s performance metric is the time to train to a defined quality target.” The choice is deliberate. That metric “captures the auxiliary operations necessary for training such models, including data-pipeline and accuracy calculations” — the very work a kernel timing excludes.

The idea came from DAWNBench, announced by the Stanford DAWN project on 29 November 2017, with a first competition deadline of 20 April 2018 and an explicit rejection of proxy metrics such as time per minibatch. The MLPerf paper credits it as “the first multi-entrant benchmark competition to use 'time to train' (originally called time to accuracy)”. MLCommons still runs it that way — “The MLPerf Training benchmark suite measures how fast systems can train models to a target quality metric” — and publishes each benchmark as a dataset plus a target. In MLPerf Training v6.0: Llama 3.1 8B to 3.3 log perplexity on C4, Llama 3.1 405B to 5.6, the Criteo 4TB recommendation task to 0.8032 AUC, lightweight object detection (RetinaNet on Open Images) to 34.0% mAP. Every entry is a finish line, not a speed.

Report valid tokens, unique examples, optimizer updates, energy, and time to a quality threshold. Energy is not a rhetorical item on that list. Two groups have published training budgets in those units.

Start with GPT-3. Patterson and colleagues, at Google and UC Berkeley, put its training at 1,287 MWh and 552 tCO2e — “Its estimated carbon emissions due to training are 552 tCO2e and its energy consumption is 1287 MWh” — and argued that energy usage and CO2e should be a key metric in evaluating models.

BLOOM-176B was measured rather than estimated. Luccioni and colleagues logged 118 days 5h 41m, 1,082,990 A100-80GB GPU hours, and 433,196 kWh drawn on a 57 gCO2eq/kWh grid: “We estimate that BLOOM’s final training emitted approximately 24.7 tonnes of CO2eq if we consider only the dynamic power consumption, and 50.5 tonnes if we account for all processes ranging from equipment manufacturing to energy-based operational consumption.” Their Table 4 reproduces the other group's figure independently, listing GPT-3 at 1,287 MWh and 502 tonnes, or 552 tonnes once a datacenter PUE of 1.1 is applied.

Two models of comparable size, roughly a twentieth of the emissions. No throughput counter would have shown that.

Optimization engineering should compare model scale, data mixture, training duration, and hardware techniques under one outcome budget. Efficiency is a design property of the whole experiment. The choice of model, datacenter and processor moves a training run's carbon footprint by ~100-1000X. No kernel rewrite reaches that range.