A GPT from Scratch

Dive into Deep Learning · §11.2

A GPT from scratch
one class, many configurations · train it, break it, sample from it · load the real GPT-2

From blocks to a language model

Four ingredients on top of the stacked block:

• Token embedding + tied head — the same table in and out (Press & Wolf, 2017); GPT-2 ties too.
• Causal mask — as a flag of the fused attention kernels; the n \times n score matrix is never materialized.
• Positions — 'learned' table added at the bottom (GPT-2) or 'rope' rotations inside every head (the modern default).
• The blocks — d2l.TransformerBlock, unchanged, via attn_factory.

The GPT class

class GPT(nn.Module):
    """Decoder-only transformer language model built from configurable
    blocks."""

    class CausalAttention(nn.Module):
        """Multi-head causal self-attention, optionally rotary."""
        def __init__(self, num_hiddens, num_heads, bias=False, rope=False):
            super().__init__()
            self.num_heads, self.rope = num_heads, rope
            self.W_qkv = nn.Linear(num_hiddens, 3 * num_hiddens, bias=bias)
            self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

        def _rope(self, x):
            d = x.shape[-1]
            pos = torch.arange(x.shape[-2], dtype=torch.float32,
                               device=x.device)
            inv_freq = 10000.0 ** (
                -torch.arange(0, d, 2, device=x.device) / d)
            theta = pos[:, None] * inv_freq[None, :]
            cos, sin = torch.cos(theta), torch.sin(theta)
            x1, x2 = x[..., 0::2], x[..., 1::2]
            return torch.stack([x1 * cos - x2 * sin,
                                x1 * sin + x2 * cos], -1).flatten(-2)

        def forward(self, X, *_):
            B, T, D = X.shape
            q, k, v = self.W_qkv(X).chunk(3, -1)
            q, k, v = (u.reshape(B, T, self.num_heads, -1).transpose(1, 2)
                       for u in (q, k, v))
            if self.rope:
                q, k = self._rope(q), self._rope(k)
            Y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
            return self.W_o(Y.transpose(1, 2).reshape(B, T, D))

    def __init__(self, vocab_size, num_hiddens=256, num_heads=8, num_blks=6,
                 max_len=1024, pos='rope', norm='rms', act='swiglu',
                 pre_norm=True, bias=False, dropout=0):
        super().__init__()
        self.pos, self.max_len = pos, max_len
        self.token_emb = nn.Embedding(vocab_size, num_hiddens)
        nn.init.normal_(self.token_emb.weight, std=0.02)
        if pos == 'learned':
            self.pos_emb = nn.Embedding(max_len, num_hiddens)
            nn.init.normal_(self.pos_emb.weight, std=0.02)
        attn = lambda: self.CausalAttention(num_hiddens, num_heads, bias,
                                            rope=(pos == 'rope'))
        self.blks = nn.ModuleList([
            d2l.TransformerBlock(num_hiddens, num_heads, dropout, norm, act,
                                 pre_norm, bias, attn_factory=attn)
            for _ in range(num_blks)])
        self.norm = (nn.RMSNorm if norm == 'rms'
                     else nn.LayerNorm)(num_hiddens)

    def forward(self, X):
        H = self.token_emb(X)
        if self.pos == 'learned':
            H = H + self.pos_emb(torch.arange(X.shape[1], device=X.device))
        for blk in self.blks:
            H = blk(H)
        return F.linear(self.norm(H), self.token_emb.weight)

Train the modern configuration

4.7M parameters, character-level Time Machine, one GPU, about a minute:

final: train 0.14, validation 3.44; best validation 1.49 at step 200

Reading the two curves

• Train falls smoothly to ~0.15 nats; validation bottoms near 1.5 within a few hundred steps, then rises.
• 16M training tokens over a 100k-character book ≈ 160 passes, with dozens of parameters per character: past the first epochs there is nothing to learn but the book itself.
• The cure is more data, not more steps — the scaling-laws section’s subject.

Cost anchor: this run ≈ 6ND \approx 5 \times 10^{14} FLOPs. GPT-2 (124M) ≈ 7 \times 10^{18}. Frontier ≈ 10^{25}+. Same block at every scale — the systems around it change.

Breaking it with one flag

pre_norm=False, learning rate 3\times10^{-3} (fine for pre-norm):

pre_norm=True: loss at step 100/200/400/800: 2.20/1.81/1.36/1.06
pre_norm=False: loss at step 100/200/400/800: 2.83/2.83/2.83/2.83

No divergence — worse: pinned at 2.83 nats = the unigram entropy of the text. Attention never learns to use context, exactly as the at-initialization gradients predicted.

Sampling: temperature and truncation

@d2l.add_to_class(GPT)
@torch.no_grad()
def generate(self, prefix, num_tokens, temperature=1.0, top_k=None):
    """Sample a continuation of the token-id list prefix."""
    ids = list(prefix)
    device = next(self.parameters()).device
    for _ in range(num_tokens):
        X = torch.tensor(ids[-self.max_len:], device=device)[None]
        logits = self(X)[0, -1] / temperature
        if top_k is not None:
            cutoff = torch.topk(logits, top_k).values[-1]
            logits[logits < cutoff] = -torch.inf
        ids.append(int(torch.multinomial(F.softmax(logits, -1), 1)))
    return ids

Deliberately naive: every token reruns the full forward pass — measuring and fixing that is the next section (KV cache).

What a minute of training sounds like

T=1.0, top_k=None: 'the time traveller put forth his finger towards the lever no he said suddenly lend me your hand and turning to the psychologist he took tha'
T=0.7, top_k=8: 'the time traveller pushed his plate away and looked round us i suppose i must apologize he said i was simply starving i ve had a most amazi'
T=2.0, top_k=None: 'the time traveller then when we had all imitated the actual asory may ejustionted you cannot move about in all of them i shook her off litt'

Fluent pseudo-Wells — much of it apparently memorized Wells: the validation gap made audible. Temperature barely matters when a memorizing model is this confident.

Loading GPT-2: config = constructor call

GPT-2 (124M) is our class with the 2019 flags: pos='learned', norm='layer', act='gelu', pre_norm=True, bias=True.

• Weights + tokenizer files pinned by sha1 in d2l.DATA_HUB.
• Tokenizer: GPT-2’s merge list + vocabulary + ch. 8’s BPE pattern, assembled into tiktoken — no model library.

HF_URL = 'https://huggingface.co/openai-community/gpt2/resolve/main/'
d2l.DATA_HUB['gpt2-weights'] = (
    HF_URL + 'model.safetensors',
    '89a76996d7c6ee89b86618a265483aab73e61d50')
d2l.DATA_HUB['gpt2-merges'] = (
    HF_URL + 'merges.txt', '396d4d8ec90cb02f4d56e049e0e4add868bcd943')
d2l.DATA_HUB['gpt2-encoder'] = (
    HF_URL + 'vocab.json', 'f0223209235343bc067d7da838328bced8085ae1')

enc = tiktoken.Encoding(
    'gpt2', explicit_n_vocab=50257,
    pat_str=d2l.BPETokenizer.GPT2_PATTERN,
    mergeable_ranks=data_gym_to_mergeable_bpe_ranks(
        d2l.download('gpt2-merges', '../data/gpt2'),
        d2l.download('gpt2-encoder', '../data/gpt2')),
    special_tokens={'<|endoftext|>': 50256})
ids = enc.encode('Attention is all you need.')
print(ids, [enc.decode([i]) for i in ids])

[8086, 1463, 318, 477, 345, 761, 13] ['Att', 'ention', ' is', ' all', ' you', ' need', '.']

The weight mapping

One dictionary from checkpoint names to modules; the one trap is GPT-2’s Conv1D layout, (\textrm{in}, \textrm{out}) — transpose for nn.Linear, adopt unchanged for nnx.Linear:

gpt2 = GPT(vocab_size=50257, num_hiddens=768, num_heads=12, num_blks=12,
           max_len=1024, pos='learned', norm='layer', act='gelu',
           pre_norm=True, bias=True)  # the GPT-2 (124M) configuration
weights = load_file(d2l.download('gpt2-weights', '../data/gpt2'))

with torch.no_grad():
    gpt2.token_emb.weight.copy_(weights['wte.weight'])
    gpt2.pos_emb.weight.copy_(weights['wpe.weight'])
    gpt2.norm.weight.copy_(weights['ln_f.weight'])
    gpt2.norm.bias.copy_(weights['ln_f.bias'])
    for i, blk in enumerate(gpt2.blks):
        modules = {f'h.{i}.ln_1': blk.norm1, f'h.{i}.ln_2': blk.norm2,
                   f'h.{i}.attn.c_attn': blk.attention.W_qkv,
                   f'h.{i}.attn.c_proj': blk.attention.W_o,
                   f'h.{i}.mlp.c_fc': blk.ffn.W_1,
                   f'h.{i}.mlp.c_proj': blk.ffn.W_2}
        for key, module in modules.items():
            W = weights[key + '.weight']
            module.weight.copy_(W.T if W.ndim == 2 else W)  # Conv1D layout
            module.bias.copy_(weights[key + '.bias'])

gpt2.to(d2l.try_gpu()).eval()
print(f'{sum(p.numel() for p in gpt2.parameters()) / 1e6:.1f}M parameters')

124.4M parameters

Did it work?

Silent-failure insurance, one number and one sentence:

per-token loss 3.89, perplexity 49

Alan Turing theorized that computers would one day become the most powerful machines on the planet.

The computer is a machine that
The secret of a good deep learning textbook is that you don't know what you're supposed to write. That's not quite right, but it's a good idea.

So let's dive into the book for a second.

Recap

• GPT = embedding + causal blocks + final norm + tied head; positions learned-or-rotary; one class, flags for a decade of designs.
• Read both curves: best validation early, then memorization — data, not steps.
• pre_norm=False at a healthy learning rate pins training at the unigram plateau.
• Sampling = temperature + truncation; naive generation recomputes everything (the KV cache fixes this, next).
• The released GPT-2 loads into our class and completes English the way the history books say it should.