Inference costs

What you will be able to do

Know before you deploy whether a model fits and how fast it answers. You estimate the memory of SmolLM2-135M from its parameter count, see what the KV cache stores and what it saves, and read real measurements of batching and quantization. Then you can explain why training needs far more memory than inference.

NotebookDo this lesson in a notebookColab · Kaggle · GitHub

The notebook computes the memory of the weights and looks inside the KV cache. It times generation on your computer, with and without the cache and with a batch. It also tries int8 quantization and counts the memory of one real training step. About 30 minutes. The last section needs about 3 GB of free memory.

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. They need a free Google or Kaggle account. Stuck? See the solution notebook.

Before you startWhat to know and to have readyAbout 55 minutes

What you need to know

• SmolLM2-135M, its 30 blocks and grouped-query attention, and the generation loop, from Training versus inference.
• Queries, keys and values, from Attention step by step.
• The dtype float32, from Tensors and shapes.
• Optional: latency, throughput and the idea of quantization, from Resource trade-offs in the language models course.

What you need on your computer

• Nothing for the notebook route.
• For the local route: the setup of Training versus inference, with SmolLM2-135M loaded as model and tokenizer, and PROMPT and show_shape() defined as there.

The course measured the timings on this page on a quiet Mac mini, with PyTorch 2.14.1 and transformers 5.19.0. It used the CPU with 4 threads and took the median of 5 runs. Your times will differ; the comparisons should not.

Re-reading a letter for every word

Imagine that you write a reply to a long letter, with a strange rule. Before each new word, you read the whole letter again from the start, and your reply so far. The first word is quick. The fiftieth word means reading everything fifty times. A sensible person reads the letter once, keeps notes, and adds one line to the notes for each word written.

Notes are not free: they take space on the desk, and the space grows with every word. Generation with a language model has the same trade-off.

The problem: will it fit, and how fast is it?

Amira must run SmolLM2-135M inside the help-desk app. Before she asks for a server, she wants three numbers. How much memory does the model need? How long does one reply take? What happens when ten agents use it at once? If she guesses, the app either crashes for lack of memory or wastes money on a server that is far too large.

Most of the answer is arithmetic that you can do before you run anything.

Memory, part 1: the weights

Each parameter is one number, and each number takes a fixed number of bytes. float32 (fp32) takes 4 bytes. bfloat16 (bf16) is a 16-bit format with the same range as float32 but fewer digits: 2 bytes. int8 takes 1 byte, and int4 half a byte. So:

Format	Bytes per parameter	SmolLM2-135M (134,515,008 parameters)
float32	4	538 MB
bfloat16	2	269 MB
int8	1	135 MB
int4	0.5	67 MB

The real tensors agree with the arithmetic. In the notebook, the loaded float32 model holds exactly 538,060,032 bytes of weights: 134,515,008 × 4.

The makers store SmolLM2’s weights in bfloat16; that is why the download is about 270 MB. The course loads them as float32, so the copy in memory is twice as large but holds exactly the same values. There are two reasons. Every CPU computes float32 well, and the course’s numbers were all measured in float32.

Memory, part 2: activations

Activations are the numbers that a forward pass computes on its way. (The word also names functions such as ReLU, from Why layers need activations; here it means the outputs of the layers.) They are the hidden states after each block, the attention scores and the logits. They exist only during the pass. In inference, each block’s activations can be freed as soon as the next block has used them. They grow with the number of texts and the number of tokens.

For the prompt of 25 tokens, one hidden state has the shape (1, 25, 576): 57,600 bytes in float32. The logits are much larger, (1, 25, 49152): 4,915,200 bytes, almost 5 MB. That is because every position gets a score for every token of the vocabulary. The attention scores of one block, 9 heads × 25 × 25, are tiny for 25 tokens. But they grow with the square of the length, so long texts make activations matter.

Memory, part 3: the KV cache

Recall attention: each token’s query is compared with the keys of the tokens before it, and the weights mix their values. With the causal mask, the keys and values of the earlier tokens never change when a new token arrives. So there is no reason to compute them again. The KV cache keeps the keys and values of every token so far, in every block, like the notes on the desk.

That splits generation into two phases. The prefill runs the whole prompt through the model once and fills the cache. Then each decode step feeds only the newest token. That token’s query reads the cached keys and values, and its own key and value are added to the cache.

The generation loop with a KV cache: one prefill for the prompt, then one cheap decode step per new token. Times from the course's measurements (Mac mini CPU, 4 threads, medians of 5).

Text description of the diagram

The top row shows the phases. First the prefill, once: all 25 prompt tokens go through the model in one forward pass, in about 0.02 seconds. Then decode step 1: the input is the token “is”, and the output is “very”, in about 0.011 seconds. Decode step 2: in “very”, out “helpful”. Decode step 3: in “helpful”, out a full stop. This repeats until 40 new tokens or an end token.

Below each phase is the KV cache. After the prefill, it holds 25 tokens: shape (1, 3, 25, 64) in each of the 30 blocks. Each decode step reads the cache and adds one token: 26, 27, then 28 tokens.

At the bottom, dashed: without the cache, every step runs the model on the whole text again: 25 tokens, then 26, 27, up to 64 tokens for 40 new tokens. The answer is the same, but the work grows with every step: 1.01 seconds instead of 0.46 seconds on the course’s CPU.

Amira looks inside. use_cache=True asks the model to return the cache, as past_key_values, after the prefill:

kv_cache.py

inputs = tokenizer(PROMPT, return_tensors="pt")
with torch.no_grad():
    out = model(**inputs, use_cache=True)              # the prefill
cache = out.past_key_values
print("blocks in the cache:", len(cache.layers))
show_shape("keys", cache.layers[0].keys)
show_shape("values", cache.layers[0].values)
cache_bytes = sum(l.keys.numel() * 4 + l.values.numel() * 4 for l in cache.layers)
print("cache bytes for 25 tokens:", cache_bytes, "| per token:", cache_bytes // 25)

The output:

blocks in the cache: 30
keys         shape (1, 3, 25, 64)
values       shape (1, 3, 25, 64)
cache bytes for 25 tokens: 1152000 | per token: 46080

Read the shape (1, 3, 25, 64) as: 1 text, 3 key-value heads, 25 tokens, 64 numbers. Then one decode step, with only the newest token as input:

kv_cache.py

next_id = out.logits[0, -1].argmax().reshape(1, 1)
with torch.no_grad():
    step = model(input_ids=next_id, past_key_values=cache, use_cache=True)
show_shape("input", next_id)
show_shape("logits", step.logits)
show_shape("keys now", cache.layers[0].keys)

The output:

input        shape (1, 1)
logits       shape (1, 1, 49152)
keys now     shape (1, 3, 26, 64)

One token in, one row of scores out, and the cache grew by one token.

How big is the cache?

Each token stores a key and a value (2), in 30 blocks, for 3 key-value heads of 64 numbers, at 4 bytes each:

2 × 30 × 3 × 64 × 4 = 46,080 bytes per token.

The prompt of 25 tokens needs 1,152,000 bytes, exactly the size of PyTorch’s real cache above. A text of 1,024 tokens needs 47 MB. At the model’s longest context, 8,192 tokens, it needs 377 MB: more than the int8 weights. And each text in use needs its own cache.

Here grouped-query attention pays off. With one key-value head for each of the 9 query heads, the cache would be three times as large: 138,240 bytes per token, 1.13 GB at 8,192 tokens.

• Weights
• KV cache

Memory of SmolLM2-135M: the weights in four number formats (parameters × bytes), and the KV cache of one text at 1,024 and 8,192 tokens in float32 (46,080 bytes per token).

Text description of the chart

A column chart in megabytes. Weights: float32 538, bfloat16 269, int8 135, int4 67. KV cache of one text in float32: 47 at 1,024 tokens and 377 at 8,192 tokens.

What the cache saves

The course measured 40 new tokens of greedy generation for the Larkfield prompt, with and without the cache. The prefill of the 25 prompt tokens took about 0.02 seconds, and one decode step about 0.011 seconds.

Time to generate 40 tokens from the 25-token Larkfield prompt with SmolLM2-135M: float32 with the KV cache, float32 without it, and int8 with the cache. Measured on the course's Mac mini CPU, 4 threads, median of 5 runs. The int8 column is explained below.

Text description of the chart

A column chart of the seconds for 40 new tokens. With the KV cache: 0.46. Without the cache: 1.01. Int8 weights with the cache: 0.60.

Without the cache, the same 40 tokens took 1.01 s instead of 0.46 s, and the tokens were identical. The gap grows with the length of the text, because each step without the cache reads every earlier token again. The cost of the cache is memory, as the chart above shows. Generation libraries use the cache by default; use_cache=False exists to compare, or to save memory when a text is very long.

Batching: many texts, one pass

A decode step for one text does little arithmetic for each weight that it reads from memory. So a CPU spends much of the step only reading the weights. A batch of texts goes through the model together: each weight is read once and used for all of them. The course’s measurement generated 20 tokens for 8 copies of the prompt:

8 prompts × 20 new tokens	Time
One after another	1.92 s
As one batch of 8	0.47 s

The batch was four times faster for the same work. The costs: texts of different lengths need padding (filler tokens, on the left for generation). Each text in a batch waits for the whole batch, and the batch needs eight caches. For Amira, batching matters when several agents ask at the same moment, not for one agent.

Quantization: memory, speed and quality

Quantization stores the weights with fewer bits: it rounds each weight to one of a small set of values. Resource trade-offs in the language models course explains the rounding. It also shows that a smaller copy is not always as good. Here, the course measured what it does to SmolLM2 on a CPU.

PyTorch’s dynamic int8 quantization (quantize_dynamic) stores every Linear layer’s weights as 8-bit integers. During the computation, it converts the numbers it needs. The results, against float32:

	float32	Dynamic int8
Weights of the Linear layers	4 bytes each	1 byte each
40 new tokens	0.46 s	0.60 s
Probability of “ is” next	0.107	0.195
Greedy text	“ is very helpful. I am looking for a new parcel. …”	“ is not answering my emails. The answer is that …”

Three lessons in one table. Int8 saves memory. It was slower on this CPU: converting the numbers back costs time, and this processor’s int8 code is not faster than its float32 code. And it changed the outputs: the probabilities moved, so the greedy text took another path from its first steps.

Even bfloat16, a much smaller change, moves the numbers. Run in bfloat16, SmolLM2 gives “ is” 0.102 instead of 0.107. Its greedy text is the same for 10 tokens, then says “a new home” where float32 says “a new parcel”.

So quantization is a three-way trade: memory against speed against quality. Measure all three on your own hardware and your own task before you choose.

Training needs far more memory than inference

Inference needs the weights, the activations of the pass in progress, and the cache. Training needs more, for every weight:

• the weight itself;
• its gradient from the backward pass, the same size as the weights;
• the optimizer’s state. Adam keeps two more numbers per weight: a running average of the gradients and a running average of their squares.

That is four numbers for every parameter. Priya checks it with one real training step on SmolLM2 in float32, next-token prediction on the prompt, and counts the bytes:

training_memory.py

model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
batch = tokenizer(PROMPT, return_tensors="pt")
loss = model(**batch, labels=batch["input_ids"]).loss   # next-token prediction on the prompt
loss.backward()
optimizer.step()

mb = lambda tensors: sum(t.numel() * t.element_size() for t in tensors) / 1e6
weights = mb(model.parameters())
gradients = mb(p.grad for p in model.parameters())
adam = mb(v for state in optimizer.state.values() for k, v in state.items() if k in ("exp_avg", "exp_avg_sq"))
print(f"weights {weights:.0f} MB | gradients {gradients:.0f} MB | Adam {adam:.0f} MB | total {weights + gradients + adam:.0f} MB")

The output:

weights 538 MB | gradients 538 MB | Adam 1076 MB | total 2152 MB

labels= makes the model compute its own next-token loss; optimizer.state holds Adam’s two running averages, exp_avg and exp_avg_sq, for each weight tensor. 2,152 MB: four times the weights. Training also keeps the activations of every block until the backward pass has used them, and they grow with the batch and the text length. Inference with the same model needs 538 MB plus a cache of a few MB.

Common mistake

A large model trained on ordinary hardware

Symptom: Tomás proposes: “Let’s train our own large language model on the office Mac. It runs SmolLM2 easily, so it can train a big model too.” His plan names a model of 70 billion parameters and a Mac with 16 GB of memory.

Cause: inference memory is not training memory. Training in float32 needs four numbers of 4 bytes per parameter: the weight, its gradient and Adam’s two states. For 70 billion parameters, that is 70,000,000,000 × 16 bytes = 1,120 GB, before any activations. That is 70 times the Mac’s memory. Even SmolLM2-135M needs 2.15 GB for one training step’s numbers, against 538 MB to run it. And the memory is not the only limit: SmolLM2’s makers used 64 large GPUs for its 2 trillion tokens of pretraining.

Fix: do the memory arithmetic before you plan: parameters × bytes × copies. Ordinary hardware can run small models and fine-tune small ones. For a large model, adapt a pretrained one, with prompting or a fine-tuning method that trains only a few of the weights. The case study Fine-tune a small transformer does that with a small model.

What it means for a real AI service: the hardware that serves a model is usually far smaller than the hardware that trained it. Budget them separately.

When each cost matters

• Weights decide whether the model fits at all. Compute them first: parameters × bytes.
• The KV cache decides how many texts, and how long, you can serve at once. It grows with every token of every active text.
• Activations matter for long texts and large batches, and much more in training.
• Batching raises the work done per second when many requests arrive together. It does not make one request faster.
• Quantization saves memory; it may or may not save time, and it changes outputs. Re-measure quality after it.

Guided practice: Amira’s memory budget

1. Compute SmolLM2-135M’s weights in bfloat16: 134,515,008 × 2 bytes. You should get about 269 MB.
2. Compute the KV cache in float32 for one text of 1,024 tokens: 46,080 × 1,024. You should get about 47 MB.
3. Amira expects up to 10 agents at once, each with a text of up to 1,024 tokens. Compute the cache for all 10.
4. Add the float32 weights (538 MB) and the 10 caches.
5. In the notebook, run section 3 and check that the cache of the 25-token prompt has 1,152,000 bytes.

Show the answer

Ten caches of 47.2 MB are 472 MB. With the float32 weights: 538 + 472 = about 1,010 MB, so about 1 GB, plus the activations of the pass in progress and the program itself. With bfloat16 weights and a bfloat16 cache, both halve: about 505 MB. A server with 2 GB of free memory is enough; check it with a real test before you buy.

Your turn: a long ticket thread

Some Larkfield tickets grow into long threads. Predict first, then compute: how large is the float32 KV cache of one text at the full context of 8,192 tokens? Then compute it again for a model with one key-value head per query head (9 instead of 3).

Show the answer

46,080 × 8,192 = 377,487,360 bytes, about 377 MB, which is 70% of the float32 weights. With 9 key-value heads, each token stores 2 × 30 × 9 × 64 × 4 = 138,240 bytes, so 1,132,462,080 bytes at 8,192 tokens: 1.13 GB, more than twice the weights. Grouped-query attention cuts the cache to a third.

See it end to endFine-tune a small transformerCase study

The fine-tuning case study measures what the pretrained model costs to train and to serve, against word counts.

Knowledge check

Check your understanding

What does the KV cache store, why are those numbers safe to reuse, and what does it cost?

Show the answer

It stores the keys and values of every token so far, in every block: for SmolLM2, shape (1, 3, tokens, 64) per block, 46,080 bytes per token in float32. They are safe to reuse because the causal mask means a new token never changes the keys and values of earlier tokens. The cost is memory, which grows with every token of every active text: 377 MB for one text of 8,192 tokens.

Check your understanding

Predict: Amira quantizes SmolLM2 to int8 to make the help-desk replies faster on the office CPU. What may happen to memory, speed and the replies?

Show the answer

The weights of the Linear layers take a quarter of the memory. The speed may not improve: on the course’s CPU, 40 tokens took 0.60 s in int8 against 0.46 s in float32. The replies change: the probability of “ is” moved from 0.107 to 0.195, and the greedy text took another path. She must measure speed and quality again on her hardware and tickets.

Check your understanding

Why does training SmolLM2 in float32 need about 2.15 GB for its numbers when inference needs about 538 MB? Name what is added.

Show the answer

Training adds a gradient for every weight (538 MB) and Adam’s two running averages for every weight (1,076 MB), so four copies: 4 × 538 MB = 2,152 MB. Training also keeps the activations of every block for the backward pass. Inference needs only the weights, the activations of the pass in progress and the KV cache.

Recap

• Weights: parameters × bytes. SmolLM2-135M: 538 MB in float32, 269 MB in bfloat16, 135 MB in int8, 67 MB in int4.
• Activations are temporary, grow with texts × tokens, and must be kept for the backward pass in training.
• The KV cache keeps keys and values so that each decode step computes only the new token: 46,080 bytes per token, 377 MB at 8,192 tokens. 40 tokens took 0.46 s with it and 1.01 s without.
• Batching 8 texts took 0.47 s instead of 1.92 s. Quantization to int8 saved memory, was slower on this CPU (0.60 s) and changed the outputs.
• Training needs about four copies of the weights (weights, gradients, Adam’s two states), plus activations. That is 2.15 GB for SmolLM2 and about 1,120 GB for a 70-billion-parameter model.

Next: save a model and its experiment so that someone else can repeat it.

Next: Save and inspect results