Training versus inference

What you will be able to do

Use a language model that someone else trained. You load the small open model SmolLM2-135M and read its parts with the transformer block of Module 6. Then you watch it write one token at a time. Then you can say what training changes and what inference only uses.

NotebookDo this lesson in a notebookColab · Kaggle · GitHub

The notebook loads SmolLM2-135M (about 270 MB to download), prints its architecture and shapes, and runs the generation loop by hand and with generate(). It also asks the model to route 100 validation tickets. About 30 minutes; every cell runs in seconds on a CPU.

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 50 minutes

What you need to know

• The transformer block (attention, add and normalize, feed-forward, add and normalize), from Transformer blocks.
• The causal mask, which lets each token see only the tokens before it, from Position and masking.
• Next-token prediction, and why pretraining at scale is out of reach for Larkfield, from Training objectives.
• model.eval() and torch.no_grad(), from Batches, epochs and validation.
• Optional: greedy choice, sampling and temperature, from Predicting continuations in the language models course.

What you need on your computer

• Nothing for the notebook route.
• For the local route: the virtual environment from Set up PyTorch, then pip install "transformers>=5,<6". The model downloads about 270 MB on first use and needs about 600 MB of memory.

Tested with Python 3.12, PyTorch 2.14.1 and transformers 5.19.0 on a Mac mini CPU. We ran the notebook in Jupyter on macOS, not in Colab or Kaggle.

A new colleague with years of school

Larkfield hires a new support agent. She did not learn to read and write English at Larkfield: she spent years at school for that. On her first day, Grace gives her a short briefing on the five teams and shows her a few tickets. From then on, she reads every ticket and answers. She uses what she knows; answering one ticket does not change her much.

A pretrained language model goes through the same three stages: a long and expensive school, a short briefing for the job, and then the daily work.

The problem: 4,700 tickets cannot teach a language

Module 6 trained a tiny transformer from scratch on Larkfield’s tickets. It reached 94.2% on the validation tickets, below the bag-of-words network (98.3%) and the machine learning course’s TF-IDF model (98.8%). It had to learn English and routing at the same time, from 4,696 short tickets. Those tickets hold about 280,000 tokens in SmolLM2’s vocabulary.

SmolLM2-135M, the model of this module, was pretrained on 2 trillion tokens, about 7 million times more text, on 64 large GPUs (model card). Larkfield cannot do that, as Training objectives showed. But Larkfield does not have to: the makers published the trained weights under the Apache-2.0 licence.

The idea: pretraining, adaptation, inference

The life of a pretrained model has three stages. Pretraining and fine-tuning are training: they change the weights.

The three stages of a pretrained model. Pretraining and fine-tuning change the weights; prompting and inference do not.

Text description of the diagram

Three boxes from left to right, joined by arrows.

1. Pretraining. The model reads 2 trillion tokens of public text and learns next-token prediction. All 134.5 million weights change. It took 64 large GPUs, according to the model card. The model’s makers do this.
2. Adaptation, with two options. Fine-tuning is more training on your own examples: the weights change a little, at far less cost than pretraining. Prompting puts instructions and examples into the text: the weights do not change, and the input gets longer.
3. Inference. The model reads one ticket at a time and runs forward passes, one token after another. The weights are frozen. The cost comes on every request, for as long as you use the model.

Pretraining and fine-tuning are training: the weights change. Prompting and inference leave the weights as they are.

Pretraining is the long school: next-token prediction on a huge amount of text, so that the model learns words, grammar and facts. Adaptation is the briefing for one job. Fine-tuning trains the pretrained weights a little more on labelled examples. Prompting changes only the input. Inference is the daily work: the word comes from “infer”, to draw a conclusion from what you already know. In inference, the network only runs forward passes, in evaluation mode, with no gradients.

Inference is not free just because the weights are fixed. Training happens once; inference happens on every ticket, every day. Inference costs measures that cost.

Meet SmolLM2-135M

Amira loads the model with Hugging Face transformers, a library that downloads published models and builds them as PyTorch networks. from_pretrained() downloads the weights once and keeps a copy on the disk.

load_smollm.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "HuggingFaceTB/SmolLM2-135M"
REVISION = "93efa2f097d58c2a74874c7e644dbc9b0cee75a2"  # the exact version of the files that the course used
tokenizer = AutoTokenizer.from_pretrained(MODEL, revision=REVISION)
model = AutoModelForCausalLM.from_pretrained(MODEL, revision=REVISION, dtype=torch.float32)
model.eval()                                   # inference: switch off training-only behaviour
PROMPT = ("Larkfield is an online shop for home and garden products. "
          "A customer wrote: my parcel is late. The support team")
print("parameters:", sum(p.numel() for p in model.parameters()))

The output:

parameters: 134515008

AutoTokenizer loads the model’s own tokenizer: SmolLM2 cuts text into subword pieces from a vocabulary of 49,152 tokens, not into the course’s lower-case words. AutoModelForCausalLM builds a causal language model, a decoder that predicts the next token with a causal mask. revision= pins one version of the files, so that the model stays the same even if its makers update it. The last lesson of this module explains why that matters.

Read the architecture with the block diagram

print(model) lists every part in order:

print(model)

The output:

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(49152, 576)
    (layers): ModuleList(
      (0-29): 30 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=576, out_features=576, bias=False)
          (k_proj): Linear(in_features=576, out_features=192, bias=False)
          (v_proj): Linear(in_features=576, out_features=192, bias=False)
          (o_proj): Linear(in_features=576, out_features=576, bias=False)
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=576, out_features=1536, bias=False)
          (up_proj): Linear(in_features=576, out_features=1536, bias=False)
          (down_proj): Linear(in_features=1536, out_features=576, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm((576,), eps=1e-05)
        (post_attention_layernorm): LlamaRMSNorm((576,), eps=1e-05)
      )
    )
    (norm): LlamaRMSNorm((576,), eps=1e-05)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=576, out_features=49152, bias=False)
)

It looks long, but you know every part. The diagram puts the parts in the order that a prompt of 25 tokens goes through them, with the shape after each one.

SmolLM2-135M read with the transformer block of Module 6, with the shape after each part for a prompt of 25 tokens and the parameters of each part (counted with PyTorch).

Text description of the diagram

From top to bottom: the token IDs have the shape (1, 25): one text of 25 tokens. The embedding embed_tokens is a table of 49,152 rows of 576 numbers; it gives the shape (1, 25, 576).

Then comes one block, repeated 30 times. Inside it come a layer normalization (input_layernorm, an RMSNorm) and self-attention with 9 query heads and 3 key-value heads of 64 numbers each. Then an addition of the input and the change, and a second normalization. Then the feed-forward layer from 576 to 1,536 and back to 576, with the SiLU activation, and a second addition. The shape is (1, 25, 576) in and the same shape out.

After the 30 blocks comes a final normalization, still (1, 25, 576). Then lm_head, a linear layer from 576 to 49,152, uses the same matrix as embed_tokens. It gives the logits, shape (1, 25, 49,152). The last row goes through softmax to the next token: “ is” with 0.107.

Parameters, counted with PyTorch: the embedding 49,152 × 576 = 28,311,552. One block 3,540,096: attention 884,736, feed-forward 2,654,208 and the norms 1,152. The 30 blocks together 106,202,880. The final norm 576. lm_head adds nothing new, because it is tied to the embedding. Total 134,515,008.

Read it from the top, as the block diagram of Transformer blocks:

• embed_tokens is the embedding table of Tokens and embeddings: one row of 576 numbers for each of the 49,152 tokens.
• LlamaDecoderLayer × 30 is the transformer block, 30 times. (“Llama” is the family of models whose design SmolLM2 copies.) self_attn is the attention, mlp the feed-forward layer, and the two LlamaRMSNorm layers are the normalizations. RMSNorm is a simpler layer normalization: it rescales each token’s numbers but does not subtract their mean. The additions of the residual connections are in the block’s code, so print cannot show them.
• SmolLM2 normalizes before attention and before the feed-forward layer, not after. Module 6’s tiny transformer does the same (norm_first=True); it keeps deep stacks stable.
• mlp has three matrices instead of two. gate_proj passes through SiLU, a smooth relative of ReLU, and multiplies the output of up_proj, number by number. down_proj brings the 1,536 numbers back to 576.
• rotary_emb adds the positions: rotary position embeddings, named in Position and masking. It rotates the queries and keys by an angle that depends on the position, so it has no weights to learn.
• lm_head turns each token’s 576 numbers into 49,152 scores, one per token of the vocabulary.

Two details are new. First, look at the shapes of k_proj and v_proj: 576 → 192, while q_proj gives 576. The model has 9 query heads of 64 numbers (9 × 64 = 576) but only 3 key-value heads (3 × 64 = 192). This is grouped-query attention: a group of 3 query heads shares each key-value head, like 9 listeners who share 3 sets of notes. It saves parameters, and above all memory during generation, as the next lesson shows.

Second, lm_head uses the same matrix as embed_tokens: tied embeddings. The table that turns a token into 576 numbers also turns 576 numbers back into a score for each token. PyTorch stores it once, so the model has 134.5 million parameters, not 162.8 million.

Go deeper

See where these parts come from, and why SmolLM2 changed each choice of the first transformer, in The original transformer, step by step, in the optional mathematics module.

Tokens, logits and the next token

Now one forward pass, with shapes printed at each step, the course’s habit. show_shape() is a small helper that prints a name and a shape; the notebooks of this module use it too.

def show_shape(name, t):
    print(f"{name:<12} shape {tuple(t.shape)}")

inputs = tokenizer(PROMPT, return_tensors="pt")
print(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]))
show_shape("input_ids", inputs["input_ids"])
with torch.no_grad():                          # inference: no record for gradients
    logits = model(**inputs).logits
show_shape("logits", logits)

The output:

['L', 'ark', 'field', 'Ġis', 'Ġan', 'Ġonline', 'Ġshop', 'Ġfor', 'Ġhome', 'Ġand', 'Ġgarden', 'Ġproducts', '.', 'ĠA', 'Ġcustomer', 'Ġwrote', ':', 'Ġmy', 'Ġparcel', 'Ġis', 'Ġlate', '.', 'ĠThe', 'Ġsupport', 'Ġteam']
input_ids    shape (1, 25)
logits       shape (1, 25, 49152)

PROMPT, from load_smollm.py, is “Larkfield is an online shop for home and garden products. A customer wrote: my parcel is late. The support team”. The tokenizer does not know “Larkfield”, so it cuts it into L, ark and field; Ġ marks a space before a piece. The logits have one row of 49,152 scores for every position, because the causal mask lets each position predict the token after it. For the next token after the whole prompt, only the last row counts. Softmax turns it into probabilities.

Predict before you move the control: SmolLM2 gives “ is” 10.7%. Among its 20 most likely tokens, that is 18.7%. What happens to that share at temperature 0.5?

Explore

SmolLM2's next token, re-weighted by the temperature

Text so far: Larkfield is an online shop for home and garden products. A customer wrote: my parcel is late. The support team▢

Chance of the top token

18.7%

Chance of any other token

81.3%

Tokens with 1% or more

20

·is18.7%

·has15.6%

·will8.2%

·at6.1%

·said5.9%

the other 1545.6%

The thin dark mark on each bar is the share at temperature 1, the model's own probabilities. The last bar is the other 15 of the 20 most likely tokens together.

At temperature 1.0, "·is" has a 18.7% chance, so another word comes in about 8 runs of 10. At 1, these are the model's own probabilities.

SmolLM2-135M's real probabilities for the next token (float32, transformers 5.19.0), shared among its 20 most likely tokens, which hold 57.3% of the probability, and re-weighted by the temperature.

Text description of the explorer

The explorer shows the text “Larkfield is an online shop for home and garden products. A customer wrote: my parcel is late. The support team”. Below it are bars for the five most likely next tokens, and one bar for the other 15 of the 20 most likely tokens. One slider sets the temperature, from 0 to 2 in steps of 0.1.

At temperature 1: is 18.7%, has 15.6%, will 8.2%, at 6.1%, said 5.9%, the other 15 together 45.6%; all 20 tokens have 1% or more. At 0.5: is 39.0%, has 27.1%, will 7.6%, at 4.1%, said 3.9%, the other 15 together 18.2%; 11 tokens have 1% or more. At 0.3: is 58.2%, has 31.8%. At 2: is 10.4%, has 9.5%, will 6.9%. At 0, is has 100%: the choice is greedy.

SmolLM2 is far less sure here than a large model is about “a refund” in the language models course. Its top token has only 10.7%, and its 20 most likely tokens together hold only 57.3%. Many continuations are good English: “ is”, “ has”, “ will”, “ replied”. Predicting continuations explains temperature and sampling in detail; this lesson looks at the loop underneath.

The generation loop, one token at a time

Generation repeats one step: run the model, choose one token, add it to the text, run again. Priya writes the loop by hand, with greedy choice (always the most likely token). She runs the model on the whole text at every step:

generate_by_hand.py

ids = inputs["input_ids"]
chosen = []
with torch.no_grad():
    for step in range(10):
        logits = model(input_ids=ids).logits[0, -1]       # scores for the next token
        next_id = int(logits.argmax())                     # greedy: the most likely token
        chosen.append(next_id)
        ids = torch.cat([ids, torch.tensor([[next_id]])], dim=1)

logits[0, -1] takes the last row, argmax() the position of the highest score, and torch.cat(..., dim=1) appends the new token to the text. The table shows each step, with the three most likely tokens:

Step	Input tokens	Three most likely next tokens	Chosen
1	25	·is 10.7%, ·has 8.9%, ·will 4.7%	·is
2	26	·very 4.8%, ·not 3.3%, ·busy 3.3%	·very
3	27	·helpful 47.9%, ·good 5.1%, ·busy 4.7%	·helpful
4	28	. 36.5%, ·and 23.6%, , 15.0%	.
5	29	·I 18.2%, (new line) 12.9%, ·The 5.3%	·I
6	30	·am 13.4%, ·have 11.7%, ·would 7.5%	·am
7	31	·looking 9.3%, ·a 5.7%, ·very 4.3%	·looking
8	32	·for 67.0%, ·forward 13.7%, ·to 11.1%	·for
9	33	·a 38.9%, ·an 6.6%, ·the 4.8%	·a
10	34	·new 4.6%, ·good 2.8%, · 2.2%	·new

The input grows by one token at each step: 25, 26, 27 … 34. Some steps are sure: after “very”, “ helpful” has 47.9%; after “looking”, “ for” has 67.0%. Others are open: after “is”, “ very” has only 4.8%. Hugging Face’s generate() runs the same loop:

with torch.no_grad():
    out = model.generate(**inputs, max_new_tokens=10, do_sample=False, pad_token_id=tokenizer.eos_token_id)
new_ids = out[0, inputs["input_ids"].shape[1]:].tolist()
print("generate():", repr(tokenizer.decode(new_ids)))
print("same tokens as the loop:", new_ids == chosen)

The output:

generate(): ' is very helpful. I am looking for a new'
same tokens as the loop: True

do_sample=False asks for greedy choice; pad_token_id names the token for filling, which SmolLM2 does not define. The tokens are identical. generate() is faster because it does not compute the earlier tokens again at every step. It keeps them in a KV cache, the subject of the next lesson.

Greedy repeats; sampling varies

Forty greedy tokens show a problem:

' is very helpful. I am looking for a new parcel. I am looking for a new parcel. I am looking for a new parcel. I am looking for a new parcel. I am looking for'

Once a sentence appears twice, the repeated text makes the same sentence even more likely, and greedy choice takes it every time. Sampling draws each token at random, with the probabilities as the chances, and the temperature re-weights them first. Three samples at temperature 0.8, with top_p=0.95 (keep the most likely tokens that hold 95% of the probability) and seeds 1, 2 and 3:

1 ' is currently at my front door, and I am still waiting for their response.\n\nI am the only one with the'
2 " is there to help you out. They've got all the right tips and tools to make it happen, and they know exactly"
3 ' is there to help you with the delivery. They are really friendly and supportive. It takes 2-3 hours to get'

No repetition, and three different stories. A seed fixes the random draws, so seed 1 gives the same text again on the same computer. None of the three answers the customer. That is the next problem.

Why a base model rambles

SmolLM2-135M is a base model: it went through pretraining only. It learned to continue text, not to follow instructions. Assistants such as chat models add instruction tuning and more training after pretraining; Training stages in the language models course explains those stages.

Priya tries prompting as the adaptation. Her prompt lists the five teams, then one validation ticket, then Team:. The real team is payment. SmolLM2 writes:

' I am writing regarding order LK-714258. The electric kettle was listed at a specific price on the checkout page, but'

It writes another ticket. So Priya stops reading the text, and reads the probabilities instead: which of the five team names does the model rate most likely after Team:? On the first 100 validation tickets, that is right for 44. With five example tickets in the prompt, one per team (few-shot prompting), it is right for 45. Without examples, it sends 73 tickets to delivery; with them, only 2. The examples changed which team it favours, not how well it reads.

The keyword rule of Module 1 gets 57.3% of the validation tickets right; TF-IDF with logistic regression gets 98.8%. A 135-million-parameter base model with a prompt is not a ticket router. To make a pretrained model good at one task, you fine-tune it on labelled examples. The case study Fine-tune a small transformer measures what that buys.

When to use a pretrained model, and what it costs

Use a pretrained model when the task needs language that your own data cannot teach: long or varied texts, few labels, or text generation. Start from the smallest model that does the job, and adapt it. Try prompting first, because it is cheap; fine-tune when prompting is not enough. Measure it against a simple baseline on the same validation set, as the routing test above did.

Do not use one when a simple model already does the job. For Larkfield’s routing, the machine learning course’s linear model is better, faster and cheaper. A pretrained model costs memory and time on every request, and its licence and its training data come from someone else. Resource trade-offs in the language models course helps you choose a model.

Guided practice: read the model

1. Open the notebook and run the setup cells. You should see parameters: 134515008.
2. Run section 1. In the printed architecture, find k_proj. You should see out_features=192.
3. Compute the attention parameters of one block by hand from the four shapes, and run the check cell. You should see Check passed.
4. Run section 2. Write down the shape of the logits before you look.
5. Run section 3 and compare the 10 tokens of the loop with generate(). You should see same tokens as the loop: True.

Show the answer

Attention of one block: q_proj 576 × 576 = 331,776, k_proj 576 × 192 = 110,592, v_proj 576 × 192 = 110,592, o_proj 576 × 576 = 331,776. Together 884,736. With 9 key-value heads, k_proj and v_proj would also be 576 × 576, and attention would have 4 × 331,776 = 1,327,104 parameters. The logits have the shape (1, 25, 49152): one text, 25 positions, one score for each token of the vocabulary.

Your turn: count the whole model

Count the 134,515,008 parameters from the parts, by hand, before you run anything. You need: the embedding (49,152 × 576), one block (attention 884,736, feed-forward 3 × 576 × 1,536, two norms of 576), 30 blocks, the final norm (576) and lm_head. Then answer Tomás: “Without tied embeddings, how many parameters would it have?”

Show the answer

Embedding: 49,152 × 576 = 28,311,552. Feed-forward: 3 × 576 × 1,536 = 2,654,208. One block: 884,736 + 2,654,208 + 2 × 576 = 3,540,096. Thirty blocks: 106,202,880. Total: 28,311,552 + 106,202,880 + 576 = 134,515,008. lm_head adds nothing, because it uses the embedding matrix. A separate lm_head would add another 576 × 49,152 = 28,311,552: 162,826,560 parameters, 21% more.

Common mistake

A base model treated as an assistant

Symptom: Tomás types “Which team should handle this ticket?” and a ticket into SmolLM2-135M. The model does not answer. It writes more ticket text, or repeats one sentence again and again. He concludes that “the model is broken”.

Cause: SmolLM2-135M is a base model. Pretraining taught it to continue text, so it continues the text it sees. Greedy choice then repeats the most likely loop. Nothing in pretraining taught it to follow an instruction.

Fix: read the probabilities of a few allowed answers instead of free text, as Priya did with the five team names. Give examples in the prompt (few-shot). Then measure on validation tickets against a baseline: here 45% against 98.8%. If that is not enough, fine-tune on labelled tickets, or choose an instruction-tuned model.

What it means for a real AI service: a base model and an instruction-tuned model behave very differently with the same prompt. Check which one you have, and judge it on many real inputs, not on one try.

See it end to endTrain a tiny language modelCase study

The tiny language model writes real samples at four temperatures with the generation loop of this lesson.

See it end to endFine-tune a small transformerCase study

The fine-tuning case study adapts a pretrained encoder to a new task, and measures what the pretraining bought.

Knowledge check

Check your understanding

Which of these change the weights of a model: pretraining, fine-tuning, prompting, inference? Explain in one sentence each.

Show the answer

Pretraining changes all the weights: the model learns next-token prediction on a huge amount of text. Fine-tuning changes the weights a little: more training on your own labelled examples. Prompting changes no weights: only the input text changes. Inference changes no weights: the model only runs forward passes, in evaluation mode and without gradients.

Check your understanding

SmolLM2 has 9 query heads but k_proj gives only 192 numbers per token. Explain why, and predict the size of k_proj if every query head had its own key-value head.

Show the answer

SmolLM2 uses grouped-query attention: 3 key-value heads of 64 numbers (3 × 64 = 192), each shared by a group of 3 query heads. With one key-value head per query head, k_proj would give 9 × 64 = 576 numbers per token: 576 × 576 = 331,776 weights instead of 576 × 192 = 110,592. The same holds for v_proj.

Check your understanding

Predict: you run the hand-written greedy loop and generate() with do_sample=False for 40 tokens. Will the texts be the same? Which one does more computation, and why?

Show the answer

The texts will be the same: both choose the most likely token at each step from the same probabilities. The hand-written loop does more computation, because it runs the model on the whole text at every step (25 tokens, then 26, and so on). generate() keeps the keys and values of the earlier tokens in a KV cache and computes only the new token at each step.

Recap

• Pretraining (huge text, all weights change) and fine-tuning (your examples, the weights change a little) are training. Prompting and inference change no weights.
• SmolLM2-135M is the transformer of Module 6: an embedding of 49,152 × 576, then 30 blocks, then an output layer tied to the embedding. Each block has attention with 9 query heads and 3 key-value heads, a feed-forward layer of 1,536, and RMSNorm before each part. In all: 134,515,008 parameters.
• Generation is a loop: forward pass, choose one token, append, repeat. The logits have one row per position, (1, 25, 49152); only the last row picks the next token.
• Greedy choice repeats itself; sampling with a temperature varies. A base model continues text: as a router with a prompt it gets 44–45 of 100 tickets right, against 98.8% for TF-IDF.

Next: what this loop costs in memory and time, and why training costs far more.

Next: Inference costs