Training and Tuning, lesson 2
Pretraining
Next-token prediction over a huge pile of text, and what the model learns from it.
Pretraining is where a model's weights get made: one plain task, repeated trillions of times. This lesson works through the task, its score and the training loop, with the numbers from a real run. You need softmax ("Logits and softmax over the vocabulary" in Transformers) and the training loop ("Updating weights with gradients" in Neural Networks).
By the end, you can:
- Turn text into inputs and targets, and count the guesses in a sequence.
- Score a guess with cross-entropy loss, and know the loss before training starts.
- Work out a run's number of updates, and its learning rate at any step.
- Explain what guessing well teaches a model, and what a base model can't do yet.
Part 1. One task: guess the next token
Turn text into inputs and targets, and count the guesses in a sequence.
Take the sentence "the cat sat on the mat". Nobody has labelled it, yet it holds five small quizzes. After "the", what comes next? After "the cat"? After "the cat sat"? The answer to each one is the next word of the text.
Example 1.1
Five quizzes from one sentence
Treat "the cat sat on the mat" as 6 tokens, one per word. What are the model's inputs and targets, and how many guesses get scored?
- Take the inputsthe cat sat on the
Every token except the last, since nothing follows it.
- Take the targetscat sat on the mat
The same text moved one place left: each position's answer is the token after it.
- Count the guesses6 − 1 = 5
One for each position that has a next token.
Answer5 guesses, one per input position.
Each position may look only at the tokens before it. The causal mask ("Masking future tokens" in Transformers) hides the rest, so no position sees its own answer, and all five guesses are scored in one forward pass. A sequence of thousands of tokens gives thousands of guesses at once.
Next-token prediction
Guessing each token from the ones before it. It's self-supervised: the labels come from the text itself, so any text at all is training data, with nobody paid to label it.
n tokens → n − 1 guesses
Data loaders often cut each sequence one token longer than the model reads, so every position gets a target. Then a batch of 4 million tokens, the size Llama 2 used, gives about 4 million guesses for every update.
Your turn
A training sequence is 4,096 tokens long. How many next-token guesses does one forward pass score?
Part 2. Scoring a guess
Score a guess with cross-entropy loss, and know the loss before training starts.
At each position the model gives a probability to every token in its vocabulary. The score looks at only one of them: the probability it gave the token that actually came next. Then it takes the negative natural log of that number, written −ln.
Example 2.1
The loss on one token
The right next token is "mat". What's the loss if the model gave it 0.9, 0.5 or 0.01? And for a model that knows nothing, spreading its guess evenly over GPT-2's 50,257 tokens?
- Confident and right−ln 0.9 = 0.11
A tiny loss.
- A coin flip−ln 0.5 = 0.69
Half its belief on the right token.
- Nearly ruled out−ln 0.01 = 4.61
It all but ruled "mat" out, and "mat" was right: a big loss.
- Knowing nothingln 50,257 = 10.82
Each token gets 1 ÷ 50,257, and −ln of that is ln 50,257.
Answer0.11, 0.69, 4.61 and 10.82.
The log punishes confident mistakes hardest: 0.5 to 0.01 costs far more than 0.9 to 0.5. Training averages this loss over every position in the batch. Raise e to that average and you get the perplexity. An average loss of 1.53 gives e^1.53 ≈ 4.6: as unsure as picking among 4.6 equally likely tokens.
Cross-entropy loss and perplexity
The cross-entropy loss is −ln of the probability given to the right token, averaged over every position. The perplexity is e raised to that loss.
loss = −ln p(right token), perplexity = e^loss
The knowing-nothing number is a free check. A fresh model with random weights spreads its guesses roughly evenly, so its first loss should sit near ln of the vocabulary size: 10.82 for GPT-2. A first loss far from that points to a bug, caught before you spend real compute.
Your turn
You start training a fresh model with a vocabulary of about 128,000 tokens, the size of Llama 3's. What should the loss be at step zero?
Part 3. The loop, at scale
Work out a run's number of updates, and its learning rate at any step.
The loop is the one from "Updating weights with gradients": guess, measure, nudge. Only the size changes. Llama 2 7B took 4 million tokens per step, about a thousand sequences of 4,096 tokens, averaged the loss over all of them, and made one update with AdamW, which gives each parameter its own step size.
Example 3.1
Llama 2's run, step by step
Llama 2 7B read 2 trillion tokens at 4 million tokens a step. Its learning rate warmed up over 2,000 steps to a peak of 3 × 10⁻⁴, then fell along a cosine curve to 10% of the peak. How many steps is that, and what's the rate halfway down the curve?
- Count the steps2 × 10¹² ÷ (4 × 10⁶) = 500,000
Each step is one update to all 7 billion parameters.
- Find the floor10% × 3 × 10⁻⁴ = 3 × 10⁻⁵
Where the rate ends, at step 500,000.
- Go halfway down3 × 10⁻⁵ + ½ × 2.7 × 10⁻⁴ = 1.65 × 10⁻⁴
Halfway along a cosine curve, the rate is midway between peak and floor. That's step 251,000, midway between 2,000 and 500,000.
Answer500,000 steps, and a rate of 1.65 × 10⁻⁴ at step 251,000.
Early gradients are wild, and big steps on a fresh model can wreck it, so the rate starts near zero and climbs. Late in the run, small steps polish what's there instead of jumping around.
Learning-rate schedule
How the step size changes over a run. Warmup followed by cosine decay is the most common shape. With progress running from 0 to 1 over the decay:
rate = floor + ½ × (peak − floor) × (1 + cos(π × progress))
Big runs still go wrong. While Google trained its 540-billion-parameter PaLM, the loss spiked about 20 times. The fix was to restart from a checkpoint about 100 steps before each spike and skip the 200 to 500 batches around it. The spikes didn't come back.
Your turn
A run warms up to a peak learning rate of 6 × 10⁻⁴, then decays along a cosine curve to 10% of the peak. What's the rate halfway through the decay?
Part 4. What falls out of guessing well
Explain what guessing well teaches a model, and what a base model can't do yet.
Nobody told the model about geography, grammar or Python. But some tokens can only be guessed well by knowing things. You can put a number on what knowing is worth.
Example 4.1
What a fact is worth
The text says "The Eiffel Tower is in the city of Paris". One model doesn't know the fact and spreads its guess evenly over 1,000 city names. Another gives "Paris" 0.9. What's each model's loss on that token?
- Without the fact−ln (1 ÷ 1,000) = 6.91
Every city gets 0.001.
- With the fact−ln 0.9 = 0.11
Confident and right.
- The difference6.91 − 0.11 = 6.80
Knowing the fact takes 6.80 off the loss on this one token.
Answer6.91 without the fact, 0.11 with it.
Text is full of tokens like that, and pushing the loss down on all of them pushes the weights to store whatever helps. "The keys were on the table, so she picked" needs grammar: "them", not "it". A line of Python needs its syntax. A review that ends "Rating:" needs its mood.
Base model
The model you have after pretraining. It continues text very well, but it doesn't answer you: ask it a question and it may carry on with more questions, as if it were writing a quiz. Teaching it to answer is the job of "Instruction tuning", later in this module.
Your turn
A model that hasn't learned a country's capital spreads its guess evenly over 200 capital cities. One that has gives the right city 0.8. What's each model's loss on that token?
Where you'll meet it
Two pretraining questions come up often. What should the loss be at step zero? About ln of the vocabulary size, and anything far from it is a bug. Why warm up the learning rate? Early gradients are large and noisy, and full-size steps on them can wreck a fresh model.
Most engineers never pretrain, but they fine-tune and read the same curves: Meta's fine-tuning of Llama 2 used a cosine schedule too. A fine-tune starts far below ln of the vocabulary, because the model already predicts text well. If it starts near that number, the weights didn't load.
import math
# The probability a model gave the right next token at each position
# of "the cat sat on the mat".
p_right = [0.02, 0.30, 0.65, 0.80, 0.15]
losses = [-math.log(p) for p in p_right]
avg = sum(losses) / len(losses)
print([round(x, 2) for x in losses])
print(f"average {avg:.2f}, perplexity {math.exp(avg):.1f}, start {math.log(50257):.2f}")
def lr(step, peak=3e-4, warmup=2_000, total=500_000):
floor = 0.1 * peak
if step < warmup:
return peak * step / warmup
progress = (step - warmup) / (total - warmup)
return floor + 0.5 * (peak - floor) * (1 + math.cos(math.pi * progress))
print([f"{lr(s):.2e}" for s in (1_000, 2_000, 251_000, 500_000)])
# [3.91, 1.2, 0.43, 0.22, 1.9]
# average 1.53, perplexity 4.6, start 10.82
# ['1.50e-04', '3.00e-04', '1.65e-04', '3.00e-05']The 0.02 on "cat" costs more than the other four positions together, which is the log punishing a confident miss. The schedule matches the figure: halfway down at step 251,000.
Check yourself
Three questions from across the lesson. Try each one before you look back.
1
Halfway through a long run, the loss suddenly jumps up and stays high. What's the usual fix?
2
After a few hundred steps, a small model's training loss is 0.05, yet the text it writes is gibberish. What's the likely bug?
3
Nobody trained a base model to translate. Yet give it "English: Good morning. French:" and it often writes "Bonjour." Where did that come from?
What to remember
1One task: guess the next token
Pretraining guesses each next token from the ones before it: a sequence of n tokens gives n − 1 scored guesses in one forward pass.
2Scoring a guess
The loss is −ln of the right token's probability, and a fresh model starts near ln of the vocabulary size: 10.82 for GPT-2's 50,257 tokens.
3The loop, at scale
Llama 2 7B took 500,000 steps of 4 million tokens, with a warmup to 3 × 10⁻⁴ and a cosine decay to 3 × 10⁻⁵.
4What falls out of guessing well
Lowering the loss forces the model to store facts and grammar: knowing the Eiffel Tower is in Paris cuts that token's loss from 6.91 to 0.11.
Recap and going further
The lesson on one page
- Every position guesses the token after it
- The labels come from the text itself
- A 4,096-token sequence gives 4,095 guesses
- Loss is −ln of the right token's probability
- A know-nothing model starts at ln(vocabulary size)
- Perplexity is e to the average loss
- Llama 2: 4 million tokens a step, 500,000 steps
- Warmup, then cosine decay to 10% of the peak
- AdamW gives each weight its own step size
- A first loss far from ln(vocabulary) means a bug
- A loss near 0 early means the answer is in view
- Loss spikes: roll back and skip the batches
- Facts, grammar, code and translation
- A base model continues text but doesn't answer
Try this
Train a tiny GPT and check its first loss
Clone nanoGPT and run its character-level Shakespeare example: python data/shakespeare_char/prepare.py, then python train.py config/train_shakespeare_char.py. Without a GPU, add the smaller settings the repo's README gives for a laptop (--device=cpu --compile=False and fewer layers). Watch the losses it prints, and sample from the model at a few checkpoints with sample.py.
You need: Python with PyTorch, or a free notebook such as Google Colab.
Look for
- Whether the first loss is close to ln 65 ≈ 4.17, since the vocabulary is 65 characters.
- How training and validation loss drift apart late in the run.
- What samples read like at a loss of 2.5, and at 1.5.
- Where
train.pywarms the learning rate up and decays it.
The words, in plain language
- Next-token prediction
- Guessing each token from the ones before it. It's self-supervised: the labels come from the text itself, so any text at all is training data, with nobody paid to label it.
- Cross-entropy loss and perplexity
- The cross-entropy loss is −ln of the probability given to the right token, averaged over every position. The perplexity is e raised to that loss.
- Learning-rate schedule
- How the step size changes over a run. Warmup followed by cosine decay is the most common shape. With progress running from 0 to 1 over the decay:
- Base model
- The model you have after pretraining. It continues text very well, but it doesn't answer you: ask it a question and it may carry on with more questions, as if it were writing a quiz. Teaching it to answer is the job of "Instruction tuning", later in this module.
Read next
- Llama 2: Open Foundation and Fine-Tuned Chat ModelsMeta, 2023Section 2 gives the pretraining recipe used in this lesson: batch size, schedule and loss curves.
- PaLM: Scaling Language Modeling with PathwaysGoogle, 2022Section 5.1 describes the loss spikes on a 540-billion-parameter run, and the roll-back-and-skip fix.
- nanoGPTAndrej Karpathy, on GitHubA short, readable pretraining loop you can run on a laptop.
Training and Tuning overviewSpotted a mistake? Tell us.