Short answer
Running a model mostly needs room for its weights, while training also needs a gradient for every weight, the optimizer's two extra copies and all the activations that must be kept for the backward pass. Measured on a 135-million-parameter model, full training took 5.1 times the memory of running it at the same batch size, and a LoRA adapter cut the peak by two-thirds.
Follow-up question How do you estimate your budget before you start a run?
The usual rule tells you if a model will fit: parameters times bytes per weight, as this calculator article explains. That number answers "can I run it?" Training is another bill entirely. People find out when a run that looked safe crashes on its first step.
This project measures that bill with one small model, SmolLM2-135M-Instruct, which has 134,515,008 parameters and an Apache 2.0 licence. Each row below ran in a new process on the CPU, using 32-bit numbers unless noted, and read the operating system's peak-memory figure. That makes it a real peak, not a guess. Peaks include some overhead from loading, and the 244 MB that Python and the libraries took before any model was already subtracted. The script is 07_memory_budget.py.
Where the memory goes
| Part | Size for this model | How it grows |
|---|---|---|
| Weights | 538 MB (Calculated: 134.5 million × 4 bytes) | With parameters and precision |
| Gradients, one per weight | 538 MB (Calculated) | With trainable parameters |
| Optimizer state, two per weight for Adam | 1,076 MB (Calculated) | With trainable parameters |
| Activations kept for the backward pass | Depends on the run | With batch × length × layers |
Running needs the first row and a smaller share of the last. Training needs all four. So before a single example is processed, training already costs four times the weights. Those three fixed rows come to 2,152 MB.
What it measured
Sequence length is 256 tokens unless the row says otherwise. Measured here on an Apple M5 Max with 64 GB, PyTorch 2.15 and Transformers 5.22.
| Job | Batch | Peak above baseline |
|---|---|---|
| Load only | 944 MB | |
| Run the model, no training | 1 | 943 MB |
| Run the model, no training | 8 | 1,269 MB |
| Full training | 1 | 3,093 MB |
| Full training | 2 | 3,267 MB |
| Full training | 4 | 4,741 MB |
| Full training | 8 | 6,484 MB |
| Full training, length 64 | 4 | 3,090 MB |
| Full training, length 512 | 4 | 6,489 MB |
| Training + gradient checkpointing | 8 | 4,102 MB |
| Training in 16-bit (bfloat16) | 8 | 4,259 MB |
| LoRA adapter + checkpointing, 460,800 trainable | 8 | 2,145 MB |
Read it this way:
- Fitting says little. At batch 8, running needs 1,269 MB and training needs 6,484 MB, 5.1 times more. Using only the weights' 538 MB as the reference, training takes 12 times as much.
- At batch 1, you are paying the fixed costs. Even a single example costs 3,093 MB. Cutting batch size has a floor, and that floor is about four times the weights.
- Length and batch weigh the same. Batch 4 at length 512 (2,048 tokens) needs about what batch 8 at length 256 needs (2,048 tokens): 6,489 and 6,484 MB. Long documents are a memory decision, so truncating or splitting them saves as much as cutting batch size.
- The fixes work on different rows. Gradient checkpointing recomputes activations instead of keeping them, which saved 37% here at the price of extra computing. 16-bit numbers shrink everything by half in principle; measured, they saved 34% because peak also includes overhead. LoRA freezes the weights and trains only small added matrices, so gradients and optimizer state almost vanish: 460,800 trainable parameters, 0.34% of the model.
Those figures are from one small model on a CPU. Graphics cards keep other buffers and were not tested here, so treat these as a floor and leave a margin.
Turning it on yourself
Here is the core of one measurement. It needs to run in a new process, because the peak figure never goes back down.
import resource, sys, torch
from transformers import AutoModelForCausalLM
def peak_mb():
v = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
return v / (1024 ** 2 if sys.platform == "darwin" else 1024)
model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct", dtype=torch.float32)
model.gradient_checkpointing_enable() # remove this line to compare
model.config.use_cache = False
model.train()
opt = torch.optim.AdamW(model.parameters(), 1e-4)
x = torch.randint(0, 1000, (8, 256))
for _ in range(2): # Adam's state exists from the second step on
model(input_ids=x, labels=x).loss.backward()
opt.step()
opt.zero_grad(set_to_none=False)
print("peak MB:", round(peak_mb()))
The second loop is not a detail. Adam allocates its state on the first update, so a test with one step reports a number that is too low and sends you into a crash on step two.
Try it: Run the script with (1, 256), (4, 256) and (8, 256) in three separate processes, plot the three peaks against the batch size, and read off your own fixed cost and cost per example.
How do you estimate your budget before you start a run?
Add up the rows, then keep a margin. A quick formula, where P is parameters in billions and B is bytes per weight:
| Job | Fixed memory, in GB |
|---|---|
| Run | P × B |
| Full training with Adam | P × B × 4 (more if the state stays at 32-bit) |
| LoRA | P × B + a very small share |
Calculated, for a model of 7 billion parameters at 2 bytes per weight, before any activations:
| Job | Fixed memory |
|---|---|
| Run | 14 GB |
| Full training | 56 GB |
| LoRA on frozen weights | About 14 GB |
So a machine with 24 GB of graphics memory can comfortably hold this model to run, cannot train all of it, and can often train an adapter on it. That last case may need the lower-precision loading explained in the quantization article, since 14 GB leaves limited room for activations. None of those three lines was run here; they are arithmetic based on the small model's pattern, and no GPU of that size was tested.
To find the activation share, do what this project does: run two real training steps with your longest example at your intended batch size, read the peak, then double the batch and see how it grows. Two minutes of that tells you more than any table. It also tells you which fix to try first: if the fixed costs dominate, use an adapter or lower precision; if activations dominate, shorten sequences, reduce batch size or turn on checkpointing. Project 13 covers what to look at before deciding that memory is the actual problem.
Download the companion code for this series (zip), with every script and the recorded results.
Sources
- Training Deep Nets with Sublinear Memory Cost — arXiv, April 21, 2016
- LoRA: Low-Rank Adaptation of Large Language Models — arXiv (Hu et al., Microsoft), June 17, 2021
- Adam: A Method for Stochastic Optimization — arXiv, December 22, 2014
From the Guide to Artificial Intelligence, by Robert Scales, published by Holistic Hybrid Inc. Share and adapt it under CC BY-NC-SA 4.0: credit the guide, link to https://guidetoartificialintelligence.com/en/learn/memory-budget-fit-vs-train/, no commercial use without permission, and keep the same licence.