Chapter 26 — Hands-on: a tiny Transformer you can run

The goal

Train a small Transformer on a text file and watch it learn to write. On a laptop CPU this takes a few minutes for something recognisably English-shaped.

Get some text

# ~1 MB of Shakespeare -- the standard toy dataset
curl -o input.txt https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt

The full model (about 80 lines)

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

# ---------------- data ----------------
text = open("input.txt", encoding="utf-8").read()
chars = sorted(set(text))
vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)

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

block_size = 128       # how many characters of context
batch_size = 32
n_embd, n_head, n_layer = 128, 4, 4
device = "cuda" if torch.cuda.is_available() else "cpu"

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

# ---------------- model ----------------
class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = nn.MultiheadAttention(n_embd, n_head, batch_first=True)
        self.ln2 = nn.LayerNorm(n_embd)
        self.ff = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd), nn.GELU(), nn.Linear(4 * n_embd, n_embd))

    def forward(self, x, mask):
        h = self.ln1(x)
        a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
        x = x + a                       # residual 1
        x = x + self.ff(self.ln2(x))    # residual 2
        return x

class TinyGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok = nn.Embedding(vocab_size, n_embd)
        self.pos = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([Block() for _ in range(n_layer)])
        self.lnf = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))
        mask = torch.triu(torch.ones(T, T, device=idx.device, dtype=torch.bool), 1)
        for b in self.blocks:
            x = b(x, mask)
        logits = self.head(self.lnf(x))
        if targets is None:
            return logits, None
        loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0):
        for _ in range(max_new_tokens):
            logits, _ = self(idx[:, -block_size:])
            logits = logits[:, -1, :] / temperature      # last position only
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, 1)
            idx = torch.cat([idx, nxt], dim=1)
        return idx

# ---------------- train ----------------
model = TinyGPT().to(device)
print("parameters:", sum(p.numel() for p in model.parameters()) / 1e6, "M")
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)

for step in range(3001):
    x, y = get_batch("train")
    _, loss = model(x, y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()

    if step % 500 == 0:
        model.eval()
        with torch.no_grad():
            vx, vy = get_batch("val")
            _, vloss = model(vx, vy)
        model.train()
        print(f"step {step:4d}  train {loss.item():.3f}  val {vloss.item():.3f}")

# ---------------- generate ----------------
start = torch.zeros((1, 1), dtype=torch.long, device=device)
print("\n" + decode(model.generate(start, 500)[0].tolist()))

What you should see

Loss ~4.2 at the start is not random. It's ln(vocab_size) — the model assigns equal probability to all ~65 characters. ln(65) = 4.17. Seeing this at step 0 is a great sanity check that your setup is correct.

What you have just built

A real, working, generative language model with a Transformer architecture, trained by gradient descent on next-token prediction. The differences between this and a frontier LLM are:

Architecturally, they are the same thing. That is the point of this chapter.