Bygg en liten GPT från grunden

Att bygga en liten GPT själv är det snabbaste sättet att förstå arkitekturen på riktigt. Delarna du redan kan var för sig sätts ihop:

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

Bygg och testa i den här ordningen — varje steg verifieras innan nästa:

1. Tokenizer (teckennivå räcker) → koda och avkoda en sträng, få tillbaka originalet.
2. Attention → kontrollera kausaliteten: ändra sista token, se att första utdatan är oförändrad.
3. Block → formen in = formen ut.
4. Modell → en slumpmässig modell ska ge loss ≈ ln(vokabstorlek).
5. Träning → överanpassa 100 tecken till loss nära noll. Går inte det finns en bugg.
6. Skala upp till hela korpusen.

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, vokab, d=128, h=4, lager=4, ctx=128):
        super().__init__()
        self.tok = nn.Embedding(vokab, d); self.pos = nn.Embedding(ctx, d)
        self.blocks = nn.ModuleList([Block(d, h, ctx) for _ in range(lager)])
        self.nf = nn.LayerNorm(d); self.head = nn.Linear(d, vokab, bias=False)
        self.head.weight = self.tok.weight                      # weight tying

    def forward(self, idx, mal=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 mal is None else F.cross_entropy(logits.view(-1, logits.size(-1)), mal.view(-1))
        return logits, loss

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

m = MiniGPT(vokab=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  ← slumpmässig modell, som väntat