Build a tiny GPT

The model's settings

You built this model in the last lesson. Here are its settings, in plain words.

• Width 32. Each token becomes a list of 32 numbers as it flows through the model.
• 4 heads. Each block runs 4 attention patterns side by side. Each head works with 32 ÷ 4 = 8 of those numbers.
• 2 layers. Two transformer blocks, one after the other.
• Context 24. The model reads at most 24 tokens at a time.

Together that's about 40,000 parameters, the numbers that training adjusts. The optimizer is AdamW from Phase 4. Its learning rate is 0.003, written 3e-3 in the code. Its weight decay of 0.01 gently pulls the weights toward zero, which helps against memorizing.

Run it

The cell below does all four steps. It takes a while, and the Run button counts the seconds. It uses the course's in-browser PyTorch subset. On a laptop with real PyTorch, the same code runs in about a second.

import math
import re
import requests
import torch
import torch.nn as nn

# 1. tokenize: words, numbers, punctuation and line breaks
text = requests.get("https://api.gradient.test/v1/corpus").json()["text"]
tokens = re.findall(r"[A-Za-z]+|[0-9]+|[^\sA-Za-z0-9]|\n", text)
vocab = sorted(set(tokens))
stoi = {t: i for i, t in enumerate(vocab)}
data = torch.tensor([stoi[t] for t in tokens])
split = int(0.9 * len(data))
train_data, val_data = data[:split], data[split:]
print(f"{len(tokens):,} tokens, vocabulary of {len(vocab)}")

# 2. batches of windows; y is x shifted by one token
def get_batch(d, block=24, batch=16):
    starts = torch.randint(0, len(d) - block - 1, (batch,)).tolist()
    x = torch.stack([d[i:i + block] for i in starts])
    y = torch.stack([d[i + 1:i + block + 1] for i in starts])
    return x, y

# 3. the model from the last lesson
class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h = h
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(d, dim=-1)
        q, k, v = (t.view(B, T, self.h, d // self.h).transpose(1, 2) for t in (q, k, v))
        s = (q @ k.transpose(-2, -1) / math.sqrt(d // self.h)).masked_fill(~torch.tril(torch.ones(T, T)).bool(), float("-inf"))
        x = x + self.proj((torch.softmax(s, dim=-1) @ v).transpose(1, 2).reshape(B, T, d))
        return x + self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self, vocab, context, d, n_layers, n_heads):
        super().__init__()
        self.context = context
        self.tok, self.pos = nn.Embedding(vocab, d), nn.Embedding(context, d)
        self.blocks = nn.Sequential(*[Block(d, n_heads) for _ in range(n_layers)])
        self.ln_f, self.head = nn.LayerNorm(d), nn.Linear(d, vocab)

    def forward(self, idx):
        x = self.tok(idx) + self.pos(torch.arange(idx.shape[1]))
        return self.head(self.ln_f(self.blocks(x)))

torch.manual_seed(0)
model = GPT(len(vocab), context=24, d=32, n_layers=2, n_heads=4)
print(f"{sum(p.numel() for p in model.parameters()):,} parameters")
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)

# 4. train
for step in range(301):
    x, y = get_batch(train_data)
    loss = nn.functional.cross_entropy(model(x).view(-1, len(vocab)), y.view(-1))
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 50 == 0:
        with torch.no_grad():
            xv, yv = get_batch(val_data, batch=64)
            val = nn.functional.cross_entropy(model(xv).view(-1, len(vocab)), yv.view(-1)).item()
        print(f"step {step:>3}  training loss {loss.item():.3f}  validation loss {val:.3f}")

# 5. generate
@torch.no_grad()
def generate(idx, n, temperature=0.8):
    for _ in range(n):
        logits = model(idx[:, -model.context:])[:, -1, :] / temperature
        idx = torch.cat([idx, torch.multinomial(torch.softmax(logits, dim=-1), 1)], dim=1)
    return idx

torch.manual_seed(1)
out = generate(torch.tensor([[stoi["\n"]]]), 70)[0].tolist()
words = " ".join(vocab[i] for i in out)
print("\n" + re.sub(r" ?\n ?", "\n", re.sub(r" ([.,])", r"\1", words)).strip())

Here's what each numbered part of the code does.

1. re.findall cuts the text into tokens. stoi (short for "string to integer") maps each token to its ID.
2. get_batch picks 16 random starting points. It cuts a window x and its targets y at each one, as in the worked example.
3. Block and GPT are the model from the last lesson. torch.manual_seed(0) fixes the random starting weights, so every run matches.
4. Each training step takes a fresh batch and measures the loss. Then loss.backward() works out the gradients and opt.step() nudges every weight. Every 50 steps it also measures the loss on 64 validation windows. It does this inside torch.no_grad(), so nothing is learned from them.
5. generate starts from a line break and adds 70 tokens, one at a time. Each time it divides the scores by a temperature of 0.8, which makes likely tokens a little more likely (the sampling lesson). Then torch.multinomial draws one token. The last line tidies the spaces before punctuation.

Reading the loss

The loss is cross-entropy, from Phase 2. It's the penalty for how unlikely the model thought the right next token was. Lower is better.

Before training, the model knows nothing. The best it can do is spread its bets evenly over every token. Say the vocabulary has $V$ tokens. Then each token, including the right one, gets a chance of 1 in $V$. The loss for that even guess is:

Here $\ln$ is the natural logarithm, which Python writes as math.log.

In plain words: a model that guesses evenly over $V$ tokens starts with a loss equal to the log of $V$. The more tokens there are to choose from, the higher the starting loss.

Try small numbers first. With 4 tokens, each gets a chance of 1/4 = 0.25. The loss is minus the log of 0.25, which is about 1.386. That's the same as the log of 4. The tiny GPT's vocabulary has 183 tokens, so it starts near the log of 183, about 5.2.

import math

for V in [4, 183]:
    p = 1 / V                     # even guess: every token gets the same chance
    print(V, round(-math.log(p), 3), round(math.log(V), 3))

In the training run, the loss starts at about 5.4, close to that 5.2, and falls below 1 within a few hundred steps.

The printout has two loss columns. Training loss comes from the batch the model has just learned from. Validation loss comes from the held-back text. If both fall together, the model is learning patterns that carry over to new text. If validation loss starts to rise while training loss keeps falling, the model is memorizing its training text. That's overfitting, from Phase 3.

import math

print(round(math.exp(-1), 2), round(1 / 183, 4))

import math
print(round(-math.log(1 / 8), 2))

What did it learn, and what didn't it?

The samples read like the catalog. You'll see sentence templates, real product types, and believable materials and features. Now and then there's a slip, such as "light and hikes".

Read the samples closely and you'll usually find three things.

• Grammar and templates, learned well. "The X Y is made from Z. It is A and B. Built for C."
• Words in the right slots. Product names, product types, materials and uses each turn up where they belong.
• Facts, learned poorly. It might write a camp mug "made from recycled ripstop nylon", or a sleeping bag made from fleece. Each phrase fits its slot. But those pairings never appear in the catalog.

That last point is the heart of hallucination (Phase 1). The model learned what a believable description sounds like. It didn't learn which features belong to which product.

It's like someone who has read a hundred restaurant menus. They could write a menu that sounds real, with prices and dish names in the right places. But the dishes might not exist.

Larger models train on far more data. They learn many more of these pairings, so they're right far more often. The mechanism is the same, though. They write what's likely given the text so far. That's usually what's true, but not always. This is why grounding (RAG and tools) and evaluation matter.

From here to an LLM

A large language model (LLM) is your tiny GPT with every number turned up. Here's the comparison.

Pretraining is the stage you ran above: learning to predict the next token from lots of text. The extra stages in the last row teach a model to follow instructions and act as an assistant.

The code you just ran is, in its structure, the same program.