LM1-T3 — Training Versus Inference

Learning objectives

After this lesson, you should be able to:

• Describe what happens in a training step and in an inference step, and what is fixed in each.
• Explain why training predicts all positions of a text in parallel and generation does not.
• Distinguish pretraining from fine-tuning by data and objective.
• Estimate training and inference compute and memory for a model from its parameter count.

The idea: learn once, use many times

Training is like studying: a student reads many examples, checks their answers against the right ones and adjusts. Inference is like taking the exam: no more adjusting, just applying what was learned. A language model does not learn from your prompt or your chat; its parameters are unchanged by inference. What looks like learning in a conversation is the model using text in its context, not changing itself.

Self-supervised training: the text is its own answer

To train a language model nobody has to label anything. Take any text, cut it into tokens, and use each token as the correct answer for what came before. With the text “the robot picks up”, the inputs are “the”, “robot”, “picks” and the targets are “robot”, “picks”, “up”: the targets are the inputs shifted by one position. Because the answers are in the text, the amount of training data is limited only by how much text can be collected. This is why it is called self-supervised learning, and Module 2 builds exactly this input-target construction.

Why training is parallel and inference is not

During training the whole text is known, so the model can predict the next token at every position in a single pass (a mask prevents each position from seeing later tokens, Lesson LM3-T4). During generation the future tokens do not exist yet, so each must be produced before the next can be predicted. This asymmetry explains why training can use hardware very efficiently on huge batches and why generation, one token per pass, is limited by memory bandwidth and latency.

Pretraining and fine-tuning

1. Pretraining. Train on a very large, general text collection with the next-token objective. The result is a base model (also called a foundation model) that is good at continuing text but does not follow instructions in any dependable way. This is the most expensive stage; the book trains a tiny one and loads released weights for a larger one.
2. Supervised fine-tuning. Continue training the base model on a smaller, curated, labelled dataset for a specific job: classifying messages (Module 6) or answering instructions (Module 7). The objective stays cross-entropy, but the data and often which tokens count in the loss change.
3. Preference tuning and other alignment methods. Further stages shape behaviour using comparisons of outputs rather than single right answers. They are outside the book’s core path; the InstructGPT paper in the references describes one approach.

The same architecture is used throughout. What changes between stages is the data, the loss and which parameters are updated (Module 8).

The mathematics: loss, gradient step and cost

One training step. For a batch of texts, the model with parameters θ produces probabilities for the next token at every position. The loss is the average negative log-probability of the true next tokens (cross-entropy, Lesson LM5-T1):

θ

all the trainable parameters of the model (weights of every layer).

N

the number of predicted positions in the batch (batch size times sequence length).

log P_θ(w_t | ...)

the log-probability the model gave to the correct token; it is 0 when the model is certain and right, and very negative when the true token was thought unlikely.

The gradient ∇_θL says how to change each parameter to reduce the loss; an optimiser moves the parameters a small step against it: θ ← θ − η ∇_θL in plain gradient descent, with learning rate η (Lesson LM5-T2 covers AdamW).

Compute. For a model with N_p parameters, a forward pass costs about 2N_p floating-point operations per token (one multiply and one add per parameter), the backward pass about twice that, so

These are standard approximations that ignore the attention computation, which adds a term that grows with context length; they are accurate to within a modest factor for the model sizes discussed here.

Memory. Inference needs the weights (and a cache of past keys and values, Lesson LM10-T5). Training with the Adam family of optimisers also keeps a gradient and two optimiser states for each parameter, so at 32-bit precision it needs roughly 4 times the memory of the weights alone, plus the activations saved for the backward pass.

A worked example: a 124-million-parameter model

The GPT-2 small architecture has about N_p = 124 × 10⁶ parameters. At 4 bytes per parameter the weights take 124 × 10⁶ × 4 = 496 MB.

If a CPU sustains 100 GFLOP/s on this workload (a plausible figure for a laptop; measure yours), one training step takes about 1.5 × 10¹² / 10¹¹ = 15 s. A thousand steps take more than four hours, so the tiny models in this track’s notebooks use far fewer parameters, and the full-size model is something you load, not train. The same model generates a token in about 248 × 10⁶ / 10¹¹ ≈ 2.5 ms of arithmetic; in practice memory transfers dominate and the time per token is longer (Module 10).

A training step and an inference step in PyTorch

import torch, torch.nn as nn, torch.nn.functional as F

V, d = 10, 8                                           # vocabulary size, embedding size
model = nn.Sequential(nn.Embedding(V, d), nn.Linear(d, V))   # a toy model: token ids -> logits
opt = torch.optim.AdamW(model.parameters(), lr=1e-2)

tokens = torch.tensor([[1, 2, 3, 4, 5, 6]])            # one text of 6 tokens
inputs, targets = tokens[:, :-1], tokens[:, 1:]        # targets = inputs shifted by one

# training: parameters change
logits = model(inputs)                                 # shape (1, 5, V): a prediction at every position
loss = F.cross_entropy(logits.flatten(0, 1), targets.flatten())
opt.zero_grad(); loss.backward(); opt.step()

# inference: parameters fixed
model.eval()
with torch.no_grad():
    next_id = model(tokens)[0, -1].argmax().item()     # prediction after the last token

The first block is training: the model sees five positions at once (inputs), the loss compares its five predictions with targets, backward() computes gradients and step() updates the parameters. The second block is inference: no_grad() tells PyTorch not to record gradients, and only the prediction at the last position is used. This toy model has 170 parameters; a GPT-2 small has 124 million, but the two blocks are the same shape of code. This snippet was executed for this lesson with PyTorch to confirm that the shapes and the loss decrease over a few steps.

Engineering implications and limitations

Do not expect a deployed model to learn from users. Serving does not change parameters. Systems that “remember” store text and put it into the next prompt (Module 11), or are retrained offline.

Budget training and serving separately. Training is a one-off (or periodic) cost dominated by arithmetic and memory for gradients; serving is a continuing cost dominated by latency, memory for weights and caches, and the number of users.

Choose the stage to match the job. Prompting a pretrained or instruction-tuned model needs no training. Fine-tuning needs labelled data and a GPU or patience. Pretraining from scratch is justified only with large data and compute, or for learning, which is why this track does it at toy scale.

Keep train and inference code consistent. Many bugs come from differences: a different tokenizer, a different prompt format, dropout left on during generation. Module 7 and Module 9 return to this.

Watch data leakage. Evaluation text that was in the training data makes the model look better than it is (Lesson LM2-T6).

For robots. A model on a robot is almost always used for inference only. Treat its weights as a versioned artefact, as you treat firmware: record which weights ran, and test again when they change.

Common misconceptions

Engineering takeaways

1. Training adjusts parameters to raise the probability of real next tokens; inference uses fixed parameters to generate.
2. Targets are the inputs shifted by one token, so pretraining needs no manual labels.
3. Training predicts all positions in parallel; generation is sequential, one token per pass.
4. Pretraining uses huge general text; fine-tuning uses smaller curated data for a task or behaviour.
5. Rule-of-thumb costs: about 6 N_p FLOPs per training token, 2 N_p per generated token; training memory is about 4 times the weights at 32-bit with Adam.

References and further reading

Chapter 1 of the book gives the same overview; the paper is optional.

• Ouyang, L. et al. (2022). Training language models to follow instructions with human feedback. arXiv:2203.02155: the pretraining, supervised fine-tuning and preference-tuning pipeline.
• Build a Large Language Model (From Scratch), Chapter 1: Understanding Large Language Models (primary book reference)
• Companion site for the book (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: how big these models are: parameters, tokens and context windows, and what the numbers mean for memory and compute.

Related content

Nothing below is required to finish this lesson.

• Before
  LM1-T2: Autoregressive Generation

  Generate text one token at a time by repeatedly predicting, choosing and appending a token.

• Lab
  LM1-L2: Trace Text to Generated Output

  Follow a prompt through tokens, context, probabilities and a generated continuation, and see how context changes predictions.