Training GPT

What changes in six minutes of training?

Take the GPT from the last lesson, feed it a million characters of Shakespeare and start training. After a few seconds (100 steps), ask it to write. It gives you this:

Let it run. Six minutes later (3,000 steps), the same code writes this:

It is not Shakespeare, but look at what is there: real words (“hath”, “day”, “now”, “thou”, “to this”), line breaks, a speaker’s name in capitals followed by a colon, lines that start with a capital letter.

Pause and predict No line of code changed between the two samples. What did?

Training is the process that changes those numbers. It is one short loop, repeated thousands of times: take some text, let the model guess each next character, measure how surprised it was, and nudge every number so it would be a little less surprised next time.

You already own every piece: gradient descent, backpropagation, cross-entropy. What is new is how they fit together for text, and the question every practitioner actually worries about: how do I know it is working?

The problem

You have all the pieces, but putting them to work on text leaves three practical questions open.

• Where are the examples? Gradient descent needs (input, correct answer) pairs. We have a text file. Nobody labelled it.
• Is it learning? A loss that falls is a good sign. But falling on what? A model can score perfectly by memorising.
• Which knobs matter? Learning rate, batch size, number of steps. Pick badly and the loss explodes, or crawls.

The problem

A GPT has around a million numbers (a real one: billions). They start random. We need good values for all of them.

The naive fix

Pay people to write (question, correct answer) pairs, as in classic supervised learning.

Why it fails

Far too slow and expensive. A language model needs billions of examples.

The idea

Let the text label itself. At every position, the correct answer is the character that comes next. Any text file is a huge pile of free examples.

The trade-off

The model learns whatever is in the text, good or bad. And with a small file it can memorise instead of learning patterns, so we must hold some text back to check.

A mental model

1. A text file becomes training pairs

You met this trick in Predicting the next token: slide the text by one and it labels itself. Here it is with a longer input.

Take the text “hello world” and a context length of 8. Cut out 8 characters as the input x. The target y is the same stretch shifted one character to the right.

Read one column: at position 4 the model has seen “hello” and the right answer is a space. Every column is its own little exam question.

The causal mask makes this cheap. Position 4 cannot see positions 5, 6, 7. So one pass through the model answers all 8 questions at once, without cheating.

One sequence of T tokens gives T training examples.

2. The loop

One trip round this loop is called a step. A batch is the handful of stretches processed together in one step: 32 of them in our code. Averaging over 32 gives a steadier direction than a single example would.

Step and epoch

A step is one update of the weights, made from one batch. An epoch is however many steps it takes to have seen the whole training text once.

For example: 1,000,000 characters of text, batches of 32 × 64 = 2,048 characters: one epoch is about 488 steps.

3. The exam with unseen questions

Before training, we cut the text in two: the first 90% to train on, the last 10% locked away. The model never takes a gradient step on the locked part. We only measure the loss there.

A teacher does the same thing by never setting the practice paper as the exam.

Validation loss

The loss on text the model has never trained on. It tells you whether the model learned patterns that carry over, or just memorised.

For example: train loss 1.2 and validation loss 1.3: healthy. Train loss 0.2 and validation loss 2.5: memorised.

Overfitting

The point where more training makes the model better on its training text and worse on everything else.

For example: Train loss keeps falling: 0.9, 0.6, 0.3. Validation loss turns around: 1.8, 1.9, 2.2.

The validation split plays the part of the exam paper. Unlike a student, the model has no intention to cheat: memorising is just the easiest way downhill when there is little data and a lot of capacity.

Long runs save the weights to disk every so often. That saved file is a checkpoint. You keep the one with the best validation loss, not necessarily the last one.

Let's see the numbers

Compute a loss by hand. Pretend the vocabulary has only three characters: e, a, x. The model has read “th” and produces these logits:

The text says “the”, so the right answer is e. The model gave it 66.5%. Loss = −ln(0.665) = 0.41, a small surprise.

Had the text said “thx”, the loss would be −ln(0.090) = 2.41. A big surprise gives a big loss, and the update that follows changes the weights more.

A batch with those two examples has loss (0.41 + 2.41) ÷ 2 = 1.41. That single averaged number is what loss.backward() starts from.

And the untrained model? All logits 0, every probability 1/3, loss = ln(3) = 1.10 no matter what the answer is. With 65 characters, as in our Shakespeare file: ln(65) = 4.17.

The math

Two formulas, both of which you have already used.

The second formula is plain gradient descent, which is what the playground above uses. The real code uses a refined version of it called AdamW. It is the same idea, step downhill, with three improvements to how big each step is:

• Momentum. Do not trust one noisy batch. Keep a running average of recent gradients and step along that. Like a heavy ball that keeps rolling through small bumps.
• A step size per parameter. Some weights get huge gradients, others tiny ones. AdamW divides each weight’s step by the typical size of its recent gradients, so every weight moves at a sensible pace.
• Weight decay (the W). Pull weights slightly toward zero at each step. A mild brake on memorising. (tiny_gpt.py decays every weight; larger recipes usually exempt biases and LayerNorm gains.)

That is why AdamW works with a learning rate like 0.0003 while our plain gradient descent wants 0.3: the two numbers are not comparable.

What you should remember

1. Text labels itself: y is x shifted by one. Every position is an example, and the causal mask lets one sequence of T tokens give T predictions in a single pass.
2. The loop: batch → forward → loss → backward → step. In code: get_batch, model(x, y), zero_grad, loss.backward(), opt.step(). The learning rate is the touchiest knob: too small crawls, too large diverges.
3. Loss is average surprise. A healthy untrained model starts at ln(vocabulary size): 4.17 for 65 characters. Check it every time: in our own file it caught confident copying, caused by large starting weights.
4. Hold text back. Validation loss is the honest score. Train loss falling while validation loss rises is overfitting: memorising, not generalising.

Where this appears in a real LLM

In this lesson

• One CPU, a 1 MB text file
• 300 to 3,000 steps, 2,048 tokens per step
• Constant learning rate 3e-4
• Overfitting is a real risk: the file is small
• No checkpoints: a crash costs two minutes

In a real system

• Thousands of GPUs, trillions of tokens of filtered web text, code and books
• Hundreds of thousands of steps, millions of tokens per step
• Learning rate warms up from zero, then decays
• Often about one epoch: most text is seen once, so classic overfitting is less of a worry than data quality
• Checkpoints every few hours; runs are restarted from them after hardware failures and loss spikes

Two back-of-envelope numbers

Memory: about 16 bytes per parameter. A common setup is AdamW with mixed precision: the maths runs in 16-bit, but a 32-bit copy of the weights is kept for the updates. Per parameter that is:

So a 7-billion-parameter model needs about 7 × 10⁹ × 16 = 112 GB for this alone, before counting the activations kept for backprop. That is more than most single GPUs hold, which is one reason training is split across many.

Compute: about 6 · N · D operations. For N parameters and D training tokens, the forward pass costs about 2 operations per parameter per token (one multiply, one add), and the backward pass about twice that. 2 + 4 = 6.

Our tiny GPT: 6 × 809,856 × (3,000 steps × 2,048 tokens) ≈ 3 × 10¹³. A 70B model on 1.4 trillion tokens: 6 × 7 × 10¹⁰ × 1.4 × 10¹² ≈ 5.9 × 10²³. The formula is the same, and the two answers are ten orders of magnitude apart.

That second run is DeepMind’s Chinchilla (2022), which was sized to be compute-optimal: about 20 training tokens per parameter gives the lowest loss for a fixed training budget. Recent models go far past that point. Meta trained Llama 3 70B (2024) on more than 15 trillion tokens, over 200 per parameter: 6 × 7 × 10¹⁰ × 1.5 × 10¹³ ≈ 6.3 × 10²⁴, about ten times Chinchilla’s compute for the same number of parameters. The extra training buys a better model at the same size, and the size is what sets the cost of every answer the model gives once it is deployed. Meta reported that its models were still improving after 15 trillion tokens.

Both are rules of thumb. The 16 bytes ignores activations and varies with the setup. 6·N·D ignores the attention scores, which add a little more at long context.

What you trained here is called a base model: it continues text. Turning it into an assistant takes further training stages, covered in From raw text to assistant. The next lesson freezes the weights and uses the model to generate text.