Build a small GPT from scratch

Building a small GPT yourself is the fastest way to understand the architecture properly. The parts you already know separately are put together:

tokens → embedding + position
   ↓
[ pre-norm → causal multi-head attention → residual
  pre-norm → MLP (SwiGLU) → residual ] × L
   ↓
final norm → unembedding → logits

Build and test in this order — each step is verified before the next:

1. The tokenizer (character level is enough) → encode and decode a string, get the original back.
2. The attention → check the causality: change the last token, see that the first output is unchanged.
3. The block → the shape in = the shape out.
4. The model → a random model should give a loss ≈ ln(the vocabulary size).
5. Training → overfit 100 characters to a loss near zero. If that is not possible there is a bug.
6. Scale up to the whole corpus.

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

class Block(nn.Module):
    def __init__(self, d, h, ctx):
        super().__init__()
        self.h, self.dh = h, d // h
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.proj = nn.Linear(d, d, bias=False)
        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)).view(1, 1, ctx, ctx))

    def forward(self, x):
        B, T, D = x.shape
        q, k, v = self.qkv(self.n1(x)).split(D, dim=2)
        q, k, v = (t.view(B, T, self.h, self.dh).transpose(1, 2) for t in (q, k, v))
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.dh)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf")).softmax(-1)
        y = (att @ v).transpose(1, 2).contiguous().view(B, T, D)
        x = x + self.proj(y)
        return x + self.mlp(self.n2(x))

class MiniGPT(nn.Module):
    def __init__(self, vocab, d=128, h=4, layers=4, ctx=128):
        super().__init__()
        self.tok = nn.Embedding(vocab, d); self.pos = nn.Embedding(ctx, d)
        self.blocks = nn.ModuleList([Block(d, h, ctx) for _ in range(layers)])
        self.nf = nn.LayerNorm(d); self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.tok.weight                      # weight tying

    def forward(self, idx, target=None):
        B, T = idx.shape
        x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))
        for b in self.blocks:
            x = b(x)
        logits = self.head(self.nf(x))
        loss = None if target is None else F.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1))
        return logits, loss

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

m = MiniGPT(vocab=65)
x = torch.randint(0, 65, (2, 32))
_, loss = m(x, x)
print(round(loss.item(), 2), round(math.log(65), 2))   # 4.19 4.17  ← a random model, as expected