Build a small transformer

Two ways into the same idea

Before You Start

The Core Idea

In this lesson you write the block in PyTorch, build the model and head around it, make it generate, and finally run the three checks before training.

1The block

We start from the inside, with the piece repeated five times. Look at the picture: query, key and value come out of one layer, then get split.

Exactly the block from the transformer block lesson, but with PyTorch layers so its gradients are computed automatically. qkv is one layer computing query, key and value together, which we then split; it is faster than three separate layers.

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


class Block(nn.Module):
    def __init__(self, d, heads, ctx):
        super().__init__()
        self.heads = heads
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
        self.register_buffer("mask", torch.tril(torch.ones(ctx, ctx, dtype=torch.bool)))

    def attention(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        q, k, v = (t.view(B, T, self.heads, C // self.heads).transpose(1, 2) for t in (q, k, v))
        att = (q @ k.transpose(-2, -1)) / (k.size(-1) ** 0.5)
        att = att.masked_fill(~self.mask[:T, :T], float("-inf")).softmax(dim=-1)
        return self.proj((att @ v).transpose(1, 2).reshape(B, T, C))

    def forward(self, x):
        x = x + self.attention(self.ln1(x))
        return x + self.mlp(self.ln2(x))

register_buffer stores the mask with the model and moves it to the GPU with it, without being a trainable parameter. And the last two lines of forward are pre-norm with the residual, literally as in the NumPy version.

2The skeleton and the head

The whole model is four lines in forward: token embedding plus position embedding, then 5 blocks, then a final norm, then a linear head to vocabulary size.

Add this class below Block in the same file.

class TinyGPT(nn.Module):
    def __init__(self, vocab, d=128, heads=4, layers=5, ctx=128):
        super().__init__()
        self.ctx = ctx
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(ctx, d)
        self.blocks = nn.Sequential(*[Block(d, heads, ctx) for _ in range(layers)])
        self.ln = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)

    def forward(self, idx, targets=None):
        T = idx.size(1)
        x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))
        logits = self.head(self.ln(self.blocks(x)))
        if targets is None:
            return logits, None
        return logits, F.cross_entropy(logits.flatten(0, 1), targets.flatten())

    @torch.no_grad()
    def generate(self, idx, n, temperature=1.0):
        for _ in range(n):
            logits, _ = self(idx[:, -self.ctx:])
            probs = (logits[:, -1] / temperature).softmax(-1)
            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
        return idx

idx[:, -self.ctx:] trims the input to the last 128 tokens, since the position table knows nothing further. And logits[:, -1] takes only the last position, the only one predicting a new token.

3Generation: predict, sample, append

The generate function above is a simple loop. It takes the last position's probabilities, samples a token, appends it to the input, and repeats.

import torch
from gpt import TinyGPT

torch.manual_seed(0)
chars = sorted(set("ابتثجحخدذرزسشصضطظعغفقكلمنهوي ."))
model = TinyGPT(len(chars))
start = torch.tensor([[chars.index("ا")]])
out = model.generate(start, 60)[0]
print("".join(chars[i] for i in out))

This code needs PyTorch on your machine, so the output shown is recorded from a real run. Install it with pip install torch and run it yourself.

Completely random letters, as expected. Keep this output: in the project you will compare it with what the same model generates after minutes of training.

Look at the picture: each turn adds one token. The temperature divides the scores before softmax, making sampling bolder or more careful.

4Three checks before training

Before any training, check three things: the parameter count matches what you computed, the output shape is (batch, tokens, vocab), and the first loss is near ln(V).

import math
import torch
from gpt import TinyGPT

torch.manual_seed(0)
vocab = 64
model = TinyGPT(vocab)
print("parameters:", sum(p.numel() for p in model.parameters()))

x = torch.randint(0, vocab, (8, 128))       # 8 sequences of 128 tokens
y = torch.randint(0, vocab, (8, 128))
logits, loss = model(x, y)
print("logits:", tuple(logits.shape), "-> (batch, tokens, vocab)")
print(f"loss before training: {loss.item():.2f}   ln(vocab) = {math.log(vocab):.2f}")

This code needs PyTorch on your machine, so the output shown is recorded from a real run. Install it with pip install torch and run it yourself.

1,024,384 parameters, about a million as computed. The shape (8, 128, 64) is right. And the first loss of 4.32 is close to ln(64) = 4.16; the small gap is because random initial weights are not perfectly even. If you see a first loss of 20 or 50, your initial weights are too large or the head is wrong.

How It Works, Step by Step: how TinyGPT generates text

Worked example

Everything we built in the earlier lessons comes together here. Open a file called gpt.py; we will write two classes in it, Block for the block and TinyGPT for the whole model, and use them in this lesson and the project.

Try it without code

Myth vs Reality

Try it yourself

Quick review cards

Check Your Understanding