Memory Math for Model Training

There is a running joke among people who work with neural networks: you are not really doing AI until you have seen the error message 'RuntimeError: CUDA out of memory'. The joke points at a real experience. Running out of memory is a common way that people discover their plans did not account for the resources a model needs. Training and running a neural network both consume memory, and a model that does not fit in the memory available to it cannot be trained or served at all.

This article is about doing that accounting before the error appears. It follows one question through several techniques: if a model is too large to adjust completely on the hardware you own, what can you change instead? Each technique turns out to be a way of reducing the number of trainable parameters, or the memory cost that comes with them.

The obvious way to adjust a pre-trained model is to let every weight move during the adjustment. Early on, models were small enough that this was practical, and this approach came to be called full finetuning. As models grew, the obvious approach began to fail for ordinary people, because the hardware needed to hold and update every weight became out of reach.

Full finetuning trains every parameter, and for a 7B model in a 16-bit format the total adds up to about 56 GB before activations are counted. That is more than most consumer GPUs provide. Memory cost is easiest to understand when you add up the pieces, and the first piece is the weights themselves.

The diagram separates the weights, which are 14 GB by direct calculation, from the rest of the total. The remaining 42 GB is the portion that covers gradients and optimizer states; the source does not itemize that portion, and it points to dedicated training-memory resources for the arithmetic. The point for planning is that a full finetuning run carries far more than the weights, so a GPU that could hold the 14 GB of weights can still fail the run.

The first workaround is partial finetuning. Instead of updating every parameter, you update only some of them and freeze the rest. A frozen parameter keeps its value throughout training. Picture a model with ten layers. If you freeze the first nine and train only the last one, the number of trainable parameters falls to about 10% of what full finetuning would need.

Freezing layers lowers memory use, but it does not solve the core problem. Partial finetuning is parameter-inefficient: it still needs many trainable parameters before it approaches the quality of full finetuning. One reported study on BERT large found that roughly 25% of the parameters had to be updated to match full finetuning on the GLUE benchmark. Performance rose only as more parameters were opened up.

The natural question follows: how can you reach performance close to full finetuning while training far fewer parameters? Techniques that answer this question are called parameter-efficient, and they are grouped under PEFT.

PEFT was introduced by Houlsby and colleagues in 2019. Their key move was to insert extra parameters into the model in carefully chosen places, and then train only those extra parameters. These extra pieces are called adapter modules. Their authors placed two adapter modules inside each transformer block of a BERT model.

The original weights of the block are left alone, and the adapters are the only pieces that receive updates. During finetuning, the original parameters stay unchanged, so the number of trainable parameters is simply the number of parameters inside the adapters. On the GLUE benchmark, this arrangement came within 0.4% of full finetuning while using only 3% of the trainable parameters.

Placement explains where the adapters sit. Now trace what happens to a single input as it moves through the network. The input passes through the frozen base computation, then through the adapter, and then continues onward. Only the adapter's values change during training, while the base values remain the same for every example.

Notice what this path implies. The adapter is an extra computation on every forward pass, so the model does more work than it did before adapters were added. The source identifies this as the main drawback: the additional layers add computational steps, which increases inference latency for the finetuned model. Adapters therefore trade lower training memory for slower inference, unless the trade is handled another way.

Trainability connects directly to the cost of an update. A parameter that is not trainable does not need to be updated, so its gradient does not need to be computed. Gradients are only needed for the parameters that will change, so marking parameters as not trainable removes work for them.

Memory use depends on how numbers are stored, and the source names two numerical formats. The first is BFloat16, a floating-point format that Google introduced as the secret to high performance on Cloud TPUs. The second is the integer format, which the source says is also called a fixed point format.

The source names these formats but does not describe how many bits either one uses, and it points to outside resources for the detailed arithmetic. This article therefore does not compare their storage sizes. What you can carry forward is the general rule used in the worked example: the bytes needed per parameter set the weight memory, and the format determines those bytes.

Memory calculation is crucial for both model training and model inference, and the two are treated as separate calculations. Inference memory concerns running a model. Training memory concerns training it. Before estimating anything, a practitioner needs to know which activity they are planning for, because the two are different problems.

The source does not walk through either arithmetic. For inference it cites Carol Chen's Transformer Inference Arithmetic (kipply's blog, March 2022). For training it cites EleutherAI's Transformer Math 101 (Anthony et al., April 2023). This article does not present a memory formula, because the source only points to those resources. What the source does establish is the structure of the problem: training and inference are separate calculations, each with its own reference material.

Finetuning is training work, so it belongs on the training side of this split. The source notes that engineers sometimes finetune even when they know it is not needed, because they want to learn how to do it. It also observes that a leader may find it hard to tell whether finetuning is needed or only wanted. That distinction matters for memory planning, because finetuning carries training-side memory costs whether or not it was needed.

The three approaches differ mainly in how many parameters they train and what that choice costs. The table gathers the points made so far. It summarizes the source's claims and is not a complete survey of every method.

PEFT also brings benefits beyond memory. The source says PEFT enables finetuning on more affordable hardware, and that PEFT methods are generally sample-efficient. Full finetuning may need tens of thousands to millions of examples before quality improves noticeably, while some PEFT methods can deliver strong results with only a few thousand examples.