class GPT(nnx.Module):
"""Decoder-only transformer language model built from configurable
blocks."""
class CausalAttention(nnx.Module):
"""Multi-head causal self-attention, optionally rotary."""
def __init__(self, num_hiddens, num_heads, bias=False, rope=False,
rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
self.num_heads, self.rope = num_heads, rope
self.W_qkv = nnx.Linear(num_hiddens, 3 * num_hiddens,
use_bias=bias, rngs=rngs)
self.W_o = nnx.Linear(num_hiddens, num_hiddens, use_bias=bias,
rngs=rngs)
def _rope(self, x):
# x: (batch, num_steps, num_heads, head_dim)
d = x.shape[-1]
pos = jnp.arange(x.shape[1], dtype=jnp.float32)
inv_freq = 10000.0 ** (-jnp.arange(0, d, 2) / d)
theta = pos[:, None] * inv_freq[None, :]
cos = jnp.cos(theta)[:, None, :] # broadcast over heads
sin = jnp.sin(theta)[:, None, :]
x1, x2 = x[..., 0::2], x[..., 1::2]
return jnp.stack([x1 * cos - x2 * sin,
x1 * sin + x2 * cos], -1).reshape(x.shape)
def __call__(self, X, *_):
B, T, D = X.shape
q, k, v = jnp.split(self.W_qkv(X), 3, axis=-1)
q, k, v = (u.reshape(B, T, self.num_heads, -1)
for u in (q, k, v))
if self.rope:
q, k = self._rope(q), self._rope(k)
Y = jax.nn.dot_product_attention(q, k, v, is_causal=True)
return self.W_o(Y.reshape(B, T, D)), None
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, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
self.pos, self.max_len = pos, max_len
init = nnx.initializers.normal(0.02)
self.token_emb = nnx.Embed(vocab_size, num_hiddens,
embedding_init=init, rngs=rngs)
if pos == 'learned':
self.pos_emb = nnx.Embed(max_len, num_hiddens,
embedding_init=init, rngs=rngs)
attn = lambda rngs: self.CausalAttention(
num_hiddens, num_heads, bias, rope=(pos == 'rope'), rngs=rngs)
self.blks = nnx.List([
d2l.TransformerBlock(num_hiddens, num_heads, dropout, norm, act,
pre_norm, bias, attn_factory=attn,
rngs=rngs)
for _ in range(num_blks)])
self.norm = (nnx.RMSNorm if norm == 'rms'
else nnx.LayerNorm)(num_hiddens, rngs=rngs)
def __call__(self, X):
H = self.token_emb(X)
if self.pos == 'learned':
H = H + self.pos_emb(jnp.arange(X.shape[1]))
for blk in self.blks:
H = blk(H)
return self.token_emb.attend(self.norm(H))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, viaattn_factory.
The GPT class
Train the modern configuration
4.7M parameters, character-level Time Machine, one GPU, about a minute:
final: train 0.15, validation 3.26; best validation 1.47 at step 200Reading 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: 1.94/1.48/1.10/0.54
pre_norm=False: loss at step 100/200/400/800: 2.83/2.83/2.83/2.83No 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)
def generate(self, prefix, num_tokens, seed=0, temperature=1.0,
top_k=None):
"""Sample a continuation of the token-id list prefix."""
total = len(prefix) + num_tokens
buf = jnp.zeros((1, total), dtype=jnp.int32)
buf = buf.at[0, :len(prefix)].set(jnp.asarray(prefix))
@nnx.jit
def logits_at(model, buf, t):
return model(buf)[0, t - 1]
ids, key = list(prefix), jax.random.key(seed)
for t in range(len(prefix), total):
logits = logits_at(self, buf, t) / temperature
if top_k is not None:
logits = jnp.where(logits < jnp.sort(logits)[-top_k],
-jnp.inf, logits)
key, sub = jax.random.split(key)
ids.append(int(jax.random.categorical(sub, logits)))
buf = buf.at[0, t].set(ids[-1])
return idsDeliberately 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 holding the lamp aloft i intend to explore time is that plain i was never more serious in my life none of us quite knew '
T=0.7, top_k=8: 'the time traveller holding the lamp aloft i intend to explore time is that plain i was never more serious in my life none of us quite knew '
T=2.0, top_k=None: 'the time traveller holding the lamp aloft i intend to explore time is that plain i was never more serious in my life none of us quite knew '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 = safetensors_numpy.load_file(
d2l.download('gpt2-weights', '../data/gpt2'))
gpt2.token_emb.embedding[...] = jnp.asarray(weights['wte.weight'])
gpt2.pos_emb.embedding[...] = jnp.asarray(weights['wpe.weight'])
gpt2.norm.scale[...] = jnp.asarray(weights['ln_f.weight'])
gpt2.norm.bias[...] = jnp.asarray(weights['ln_f.bias'])
for i, blk in enumerate(gpt2.blks):
for norm, key in ((blk.norm1, f'h.{i}.ln_1'),
(blk.norm2, f'h.{i}.ln_2')):
norm.scale[...] = jnp.asarray(weights[key + '.weight'])
norm.bias[...] = jnp.asarray(weights[key + '.bias'])
for linear, key in ((blk.attention.W_qkv, f'h.{i}.attn.c_attn'),
(blk.attention.W_o, f'h.{i}.attn.c_proj'),
(blk.ffn.W_1, f'h.{i}.mlp.c_fc'),
(blk.ffn.W_2, f'h.{i}.mlp.c_proj')):
# Conv1D stores (in, out) -- exactly nnx.Linear's kernel layout
linear.kernel[...] = jnp.asarray(weights[key + '.weight'])
linear.bias[...] = jnp.asarray(weights[key + '.bias'])
gpt2.eval()
n = sum(p.size for p in jax.tree.leaves(nnx.state(gpt2, nnx.Param)))
print(f'{n / 1e6:.1f}M parameters')124.4M parametersDid it work?
Silent-failure insurance, one number and one sentence:
per-token loss 3.89, perplexity 49Alan Turing theorized that computers would one day become the most powerful machines on the planet.
The computer is a computer that
The secret of a good deep learning textbook is to make sure that you have the right skills set to get started.
A good deep learning textbook is to make sure you have the right skills set to get started.
I used toRecap
- 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=Falseat 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.