AI EngineerML EngineerResearch24h

Build your own GPT

A decoder-only transformer you wrote yourself, trained on text you chose, generating new text — no framework magic, about 300 lines.

Every explanation of attention you have read describes the same equations, and reading them is not the same as having made them work. This is the build that converts that. It is also the single most common technical request in ML and research interviews — "implement causal multi-head attention" — and having typed it once, with the shape bugs, is the difference between reciting it and deriving it under pressure.

Have these behind you first

You can start without them and look things up as you go. It will take longer and you will understand less of what you typed.

PythonPyTorchNeural networksBackpropagationTransformers & attention
Stage 1 of 7, 2h

Turn text into numbers

You build: A tokenizer that encodes a string to integers and decodes them back exactly.

If a line here loses you, this is where it is taught

Start at character level. It is the wrong choice for a real model and the right choice for this one: it takes ten lines, it has no failure modes, and it lets you get to attention today rather than on Thursday.

text = open("input.txt").read()
chars = sorted(set(text))
vocab_size = len(chars)

stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for ch, i in stoi.items()}

encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)

assert decode(encode(text)) == text
Character level. Vocabulary is usually 65-100 for English prose.

Do not start with BPE

Byte pair encoding is what real models use and it is a project of its own. Build it once the model trains, from the byte pair encoding section of the tokenization topic, which trains merges on a corpus and encodes and decodes with them. Starting there is how people spend a weekend on tokenization and never reach the transformer.

Pick a corpus you will enjoy reading samples from — a few megabytes of plain text is plenty. Shakespeare is traditional because the output is obviously wrong in an entertaining way, which makes early progress visible.

Done when

decode(encode(text)) == text for your whole corpus, and you can say what your vocabulary size is and why.

Stage 2 of 7, 2h

Feed it data

You build: A function returning random batches of (input, target) pairs, offset by one token.

If a line here loses you, this is where it is taught

The entire training objective is: given some tokens, predict the next one. So a training example is a chunk of text and the same chunk shifted one position left. Everything else in this project is machinery for doing that well.

import torch

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]

block_size = 256   # how many tokens of context
batch_size = 64    # how many sequences at once

def get_batch(split):
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - block_size, (batch_size,))
    x = torch.stack([d[i : i + block_size] for i in ix])
    y = torch.stack([d[i + 1 : i + block_size + 1] for i in ix])
    return x, y

x, y = get_batch("train")
print(x.shape, y.shape)   # (64, 256) (64, 256)

Hold on to the two letters B and T. Nearly every tensor from here on is (B, T, something): batch, time, and a feature dimension. When a shape error appears — and it will — the first question is which of those three you have transposed.

Done when

You can print one batch and explain, out loud, what every dimension of its shape means and why the targets are the inputs shifted by one.

Stage 3 of 7, 6h

One head of self-attention

You build: A single attention head: queries, keys, values, a causal mask, and a weighted sum.

If a line here loses you, this is where it is taught

This is the stage the project exists for. Take it slowly — six hours is not generous, it is realistic, and most of them go on shapes.

Every position produces three vectors. A query: what am I looking for. A key: what I offer. A value: what I pass on if chosen. Score every query against every key, turn the scores into weights, and take the weighted sum of the values.

import torch.nn as nn
from torch.nn import functional as F

class Head(nn.Module):
    def __init__(self, n_embd, head_size, block_size, dropout=0.2):
        super().__init__()
        self.key   = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        # Not a parameter: a constant the optimiser must never touch.
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)      # (B, T, head_size)
        q = self.query(x)    # (B, T, head_size)

        # Scores: how much each position should attend to each other position.
        # Scale by sqrt(head_size) or softmax saturates and gradients vanish.
        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5   # (B, T, T)

        # Causal mask: position t may not see anything after t.
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)

        v = self.value(x)
        return wei @ v       # (B, T, head_size)

The mask is the whole reason this is a GPT

Remove the masked_fill and you have BERT: every position sees the whole sequence. Keep it and each position sees only what came before, which is what makes the model able to generate. If your model trains to a suspiciously low loss, you have almost certainly broken the mask and it is reading the answer.

Before moving on, print the weight matrix for one example and look at it. Check that the upper triangle is zero and every row sums to one. If you cannot see that, you have not finished this stage.

Done when

Your attention weights for each position sum to 1, and row i has exactly zero weight on every column after i. Assert both.

Stage 4 of 7, 5h

Many heads, and the block around them

You build: Multi-head attention, a feed-forward layer, residual connections and layer norm, assembled into a transformer block.

If a line here loses you, this is where it is taught

One head learns one kind of relationship. Several heads, each with a smaller dimension, learn several at once and cost the same. Run them in parallel and concatenate.

class MultiHead(nn.Module):
    def __init__(self, n_head, n_embd, block_size, dropout=0.2):
        super().__init__()
        head_size = n_embd // n_head
        self.heads = nn.ModuleList(
            Head(n_embd, head_size, block_size, dropout) for _ in range(n_head)
        )
        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        return self.dropout(self.proj(out))


class FeedForward(nn.Module):
    """Where the model thinks about what attention gathered. Four times wider
    by convention, and the convention is worth following."""
    def __init__(self, n_embd, dropout=0.2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout=0.2):
        super().__init__()
        self.sa = MultiHead(n_head, n_embd, block_size, dropout)
        self.ff = FeedForward(n_embd, dropout)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        # Pre-norm, and x + ... rather than x = ...: the residual stream is the
        # path the gradient takes back through a deep stack.
        x = x + self.sa(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x

Pre-norm, not post-norm

The original paper normalised after the sublayer. Everything since normalises before it, because pre-norm trains without a warmup schedule and post-norm often will not. If you copy the 2017 diagram exactly, expect a loss that refuses to move.

Done when

A block takes (B, T, n_embd) and returns the same shape, and stacking six of them still trains rather than producing NaN.

Stage 5 of 7, 3h

The whole model

You build: Token and position embeddings, a stack of blocks, a final norm and a projection back to vocabulary size.

If a line here loses you, this is where it is taught

Attention has no idea what order anything is in — it sees a set, not a sequence. Position information has to be added explicitly, and a learned embedding per position is the simplest thing that works.

class GPT(nn.Module):
    def __init__(self, vocab_size, n_embd=384, n_head=6, n_layer=6,
                 block_size=256, dropout=0.2):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(
            *[Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)]
        )
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok = self.tok_emb(idx)                                   # (B, T, n_embd)
        pos = self.pos_emb(torch.arange(T, device=idx.device))    # (T, n_embd)
        x = self.ln_f(self.blocks(tok + pos))
        logits = self.head(x)                                     # (B, T, vocab)

        if targets is None:
            return logits, None

        # cross_entropy wants (N, C) and (N), so flatten batch and time together.
        loss = F.cross_entropy(
            logits.view(B * T, -1), targets.reshape(B * T)
        )
        return logits, loss

The loss you should see before training

An untrained model should be guessing uniformly, so its loss is about ln(vocab_size) — roughly 4.2 for a 65-character vocabulary. Much lower means something is leaking; much higher means initialisation is broken. This one number catches more bugs in this project than any other check.

Done when

An untrained model produces a loss close to ln(vocab_size). If it does not, something is initialised wrong — find it before you train.

Stage 6 of 7, 4h

Train it

You build: A training loop with an optimiser, periodic validation loss, and a curve you can read.

If a line here loses you, this is where it is taught

model = GPT(vocab_size).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)

@torch.no_grad()
def estimate_loss(iters=200):
    model.eval()
    out = {}
    for split in ("train", "val"):
        losses = torch.zeros(iters)
        for k in range(iters):
            x, y = get_batch(split)
            _, loss = model(x.to(device), y.to(device))
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out

for step in range(5000):
    if step % 500 == 0:
        losses = estimate_loss()
        print(f"{step}: train {losses['train']:.3f}  val {losses['val']:.3f}")

    x, y = get_batch("train")
    _, loss = model(x.to(device), y.to(device))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()

Estimate the loss over many batches rather than printing the loss of a single one. A single batch is noisy enough to hide a real trend, and watching a jittery number is how people conclude training is broken when it is working.

Overfit one batch first

Before the real run, train on a single batch and confirm the loss goes to nearly zero. A model that cannot memorise one batch has a bug, and you will find it in minutes rather than after an hour of full training.

Done when

Training loss falls well below the initial ln(vocab_size), validation loss follows it down and then flattens, and you can point at where overfitting starts.

Stage 7 of 7, 2h

Make it write

You build: A generation loop that samples one token at a time and feeds it back in.

If a line here loses you, this is where it is taught

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None):
    for _ in range(max_new_tokens):
        # Never feed more context than the model has positions for.
        idx_cond = idx[:, -model.block_size:]
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :] / temperature   # only the last position

        if top_k is not None:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = float("-inf")

        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        idx = torch.cat((idx, next_id), dim=1)
    return idx

start = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(generate(model, start, 500)[0].tolist()))

Keep a sample from before training. Putting gibberish next to the trained output is the most convincing single artefact this project produces, and it costs one print statement.

This is the slow way, deliberately

Regenerating the whole context every token is O(n squared) and real inference does not do it. That is what a KV cache fixes, and implementing one on top of this is an excellent follow-up — and a common interview question in its own right.

Done when

Your model produces text in the shape of your corpus, and changing temperature visibly changes how strange the output is.

What you have at the end

A repository containing a tokenizer, a transformer you can explain line by line, a training script that produces a loss curve, and samples from before and after training. You can honestly say you have implemented a language model from scratch, and answer follow-up questions about any line of it, which is the part that actually counts.

The other projects