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 a Pre-Norm RMSNorm Configuration
4.7M parameters, character-level Time Machine, one GPU, about a minute:
final: train 0.15, validation 3.19; best validation 1.46 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. The divergence is consistent with memorization under repeated corpus reuse.
- The scaling-laws section studies how larger and more varied data changes this regime; the present run does not isolate a remedy.
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.88/1.41/1.01/0.35
pre_norm=False: loss at step 100/200/400/800: 2.83/2.83/2.83/2.83The loss remains pinned at 2.83 nats, the unigram entropy of the text. Attention does not learn to use context, consistent with the initialization-time gradient analysis.
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).
Samples after the Short Training Run
T=1.0, top_k=None: 'the time traveller hated to have servants waiting at
dinner for a hot plate at that the editor turned to his knife and fork
with a grunt an'
T=0.7, top_k=8: 'the time traveller stood before us i gave a cry of surprise
good heavens man what s the matter cried the medical man who saw him
next and t'
T=2.0, top_k=None: 'the time traveller hated to have seubitating across the
lower round nature azout you like to the buildings with intricate
parapets and type'The samples resemble the training text closely, consistent with the widening train–validation gap. Temperature has little effect because the fitted token distribution is already concentrated.
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 + our earlier 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 maps checkpoint names to modules. GPT-2’s Conv1D layout requires care: (\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 parametersValidating the Loaded GPT-2 Checkpoint
Perplexity and a completion check detect common mapping or tokenizer errors:
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.
- The validation loss reaches its minimum early and then rises as training loss falls, consistent with memorization under repeated corpus reuse.
pre_norm=Falseat a healthy learning rate pins training at the unigram plateau.- Temperature and truncation control sampling; naive generation recomputes the prefix, whereas the next section’s KV cache reuses it.
- The released GPT-2 loads into our class and produces coherent English completions for the example prompts.