Build GPT

Can you trace one token from text to prediction?

About 50 minutes · open lesson · LLM engineering

Can you follow one token through a GPT you assembled yourself?

Type “What is a cat?” into a GPT, and a fraction of a second later it has chosen the next token. In between, the text goes through a fixed sequence of steps. The tokenizer turns it into token ids. The embedding table turns each id into a vector. A stack of Transformer blocks lets those vectors read from each other through masked, multi-head attention. Out of the other end comes the next token, which the generation loop appends to the text before going round again.

You have built each of those steps separately. Joined together, they make a GPT with the same wiring as GPT-2, in about 100 lines of model code. Two exact questions are worth answering about the result.

First, what happens to one token, from the moment it is typed to the moment the next one is chosen? At every step you should be able to say what shape the data has. Second, how many parameters does the whole model have, and where do they sit? For your tiny GPT and for GPT-2 small you will be able to count them exactly, by hand.

Pause and predict After the last block, every token is a context-aware vector. Is that the next token yet?

What is still missing?

A stack of blocks takes (T, D) and returns (T, D): one context-aware vector per token. That is not yet a prediction.

To generate text you need to know which character comes next. Two small gaps remain between the output of the blocks and that answer.

Pause and predict The last block gives D numbers per position. How do you get a score for every token in the vocabulary?

The second gap is the loop around the model. A GPT predicts one token. To write a sentence, you sample a token, append it to the input, and run the whole model again.

You built that loop in Predicting the next token. Here it gets a real model inside.

One practical worry: this lesson’s code is in PyTorch, where earlier lessons used NumPy. Is the framework hiding some extra machinery?

What PyTorch does for you: exactly two things.

During the forward pass it records every operation. You did this by hand in Backpropagation, when you kept the intermediate values for the backward pass.
When you call loss.backward(), it replays your backward rules through that recording, in reverse. That is the chain rule you applied by hand.

A torch.Tensor is a NumPy array that carries that recorder. Everything else in the file, every matrix multiply and every shape, is something you have already written yourself. (It can also run the maths on a GPU, which changes the speed and leaves the maths alone.)

A mental model

When you wire up a model, draw the data rather than the boxes. Here is the whole machine in one picture, with the data drawn as grids, one row per token.

Follow the numbers 1 to 9 and watch the shape at each stage. Every stage is a lesson you have done.

Pause and predict Compare the grid that leaves the last block with the grid that went into the first one. It is…

The head (language-model head): The last layer: one matrix multiply that turns a token’s D-number vector into V raw scores, one for every token in the vocabulary.; For example: D = 128, V = 65: the vector [0.3, −1.2, …] (128 numbers) becomes 65 scores, such as “e”: 2.1, “t”: 1.7, “z”: −3.0.

A GPT is these 6 numbers

The top of the file holds the entire configuration of the model. In its wiring, GPT-2 is the same code with bigger numbers.

Five of them fix the architecture. The sixth, dropout, only matters during training (it randomly zeroes some numbers so the model cannot lean on any single one).

Every module, and where you built it

in `tiny_gpt.py`	what it is	you built it in
tok_emb	embedding lookup	Embeddings
pos_emb	a vector per slot	The Transformer block
CausalSelfAttention	masked multi-head attention	Attention, Masks and heads
FeedForward	the MLP	Neurons and layers
Block, LayerNorm, x + …	residuals and normalisation	The Transformer block
F.softmax, F.cross_entropy	scores → probabilities, and the loss	Softmax, Next token
loss.backward()	backpropagation, automated	Backpropagation
opt.step()	the training loop	Gradient descent
generate()	the autoregressive loop	Next token

Assembling a GPT is like wiring up services you have already written and tested: the new work is checking that each output shape matches the next input. Unlike services, these have no hand-written logic. Each is a bag of numbers, and all of them are tuned jointly by one loss at the very end.

Let's see the numbers

Count the repository’s tiny GPT by hand: V = 65 characters, context 64, D = 128, 4 blocks.

Pause and predict Before the table: which part of this tiny GPT holds the most parameters?

The math

One line for the whole model, reading inside-out:

For any model of serious size you can forget everything except 12 · N · D². GPT-2 small: 12 × 12 × 768² ≈ 85M of its 124M.

Weight tying: why can the head reuse the embedding table?

The embedding table E has shape (V, D): row i is the vector for token i. The head needs shape (D, V): column i produces the score for token i. Transpose E and the shapes match.

It also makes sense. With E^T as the head, the logit for token i is the dot product between the model’s output vector and token i’s embedding: “how much does what I am about to say point in the direction of this token?” Tokens with similar embeddings automatically get similar scores.

In code it is one line: self.head.weight = self.tok_emb.weight. Both names now refer to the same tensor, and gradients from both uses are added up. GPT-2 does this; it saves 38.6M parameters there. For our character-level model it saves only 8,320 (1%), because V is tiny.

It is a design choice. Several recent small models tie their weights, while many large ones keep a separate head, since V·D is a small fraction of their total anyway.

What you should remember

GPT = tokenizer → token + position embeddings → N Transformer blocks → final LayerNorm → linear head → logits → softmax → sample → append → repeat. Shapes: (T) → (T, D) → … → (T, D) → (T, V), and generation uses only the last row.
A GPT is a handful of config numbers. GPT-2 is the same code with bigger ones. Parameters ≈ 12 · N · D² plus embeddings; most live in the MLPs.
The model can only see context_len tokens. Older tokens are cropped away and have no influence.
PyTorch automates two things you did by hand: recording the forward pass and replaying the backward rules. An untrained GPT is a correct pipeline filled with the wrong numbers.

Where this appears in a real LLM

How far is our model from the real ones? In structure it is close to them. In size it is very far away.

	our tiny GPT	GPT-2 small (2019)	Llama 3 70B (dense, 2024)	DeepSeek-V3 (mixture-of-experts, December 2024)
blocks (n_layer)	4	12	80	61
vector width (n_embd)	128	768	8,192	7,168
attention heads	4	12	64 (sharing 8 key/value heads)	128 (keys and values rebuilt from one cached vector of 512 numbers)
context window	64	1,024	8,192 (128K in version 3.1)	128K
parameters	0.81M	124M	≈ 70,000M, all used for every token	≈ 671,000M in total, ≈ 37,000M used per token
tokenizer	65 characters	BPE, 50,257 tokens	BPE, ≈ 128,000 tokens	BPE, 129,280 tokens
positions / norm / MLP	learned table / LayerNorm / GELU	RoPE / RMSNorm / SwiGLU	RoPE / RMSNorm / SwiGLU, split into 256 small experts plus 1 shared, 8 of the 256 used per token

In this lesson

Character-level, trained on 1 MB of Shakespeare
Runs on a laptop CPU in minutes
Every line readable: about 100 lines of model code
The browser tracer: untrained, D = 8, 2 blocks

In a real system

Sub-word tokens, trained on trillions of tokens
Thousands of GPUs for weeks; the model is split across many devices
The same pipeline: embed, N blocks, norm, head, softmax, sample
Changed details inside the boxes: RoPE, RMSNorm, SwiGLU, grouped-query attention, no biases

GPT-2’s published architecture is what tiny_gpt.py implements: learned positions, pre-norm blocks, GELU MLPs, tied head. With vocab 50,257, context 1,024, D = 768 and 12 layers, the formula in this lesson gives exactly 124,439,808 parameters, the “124M” on the label.

Below is GPT-2 small itself, with the weights OpenAI trained, running in your browser. Every stage is a box you built in this part, and each one names the lesson that built it. It also has a page of its own, with more room.

It has the same wiring as tiny_gpt.py. The difference is the numbers inside, and those come only from training.

One difference: the wiring matches, but the starting numbers do not. tiny_gpt.py keeps PyTorch’s defaults, which fill the embedding table with numbers of size about 1. GPT-2 starts every weight small (standard deviation 0.02) and shrinks the layers that write into the residual stream by a further 1/√(2N) for N blocks. You will see in the next lesson why that choice matters on the very first step.

Before its first training step, tiny_gpt.py reports params 0.81M and writes a line of gibberish. The pipeline is correct and the numbers in it are wrong. The next lesson is about finding the right numbers.

Open the interactive lesson

The interactive lesson adds the predictions with their answers, the labs, and graded practice with spaced review.