The posts so far built a Transformer one piece at a time: the embedding table (Embeddings and Contrastive Learning), cross-entropy on the next token (Cross-Entropy and the Softmax Gradient and What “Generative” Means), AdamW and warmup (Initialization and Optimizers), normalization (Dropout and Normalization), attention (The Transformer’s Three Attention Patterns), position (Where Am I? Position in Attention, from Sinusoids to RoPE), and the residual stream holding it together (The Residual Stream). This post puts them together into one model, a small GPT, trains it on this blog, and answers two questions with measurements:
- Does it beat what came before? What “Generative” Means fixed the target in advance. On the same two held-out posts, the character LSTM scored 2.66 bits per byte, a Kneser–Ney counting model 2.57, and bz2, a file compressor given the training posts, 2.34. Lower is better.
- What is each piece worth? Take one out, retrain, and measure.
Everything here ran on a laptop GPU (an Apple M4, through PyTorch’s mps
backend), in 10 to 15 minutes per training run, and about 314 minutes for
all of this post’s code, about thirty runs. The code picks CUDA, then Apple’s GPU, then
the CPU.
The whole model on one page: every box is a piece built in an earlier post, and the green line is the residual stream that ties them together.
- A 1.8-million-parameter GPT in 80 lines: embeddings, RoPE, RMSNorm, causal attention, a SwiGLU MLP and residual connections, the Llama layout at a thousandth of the size.
- On the two held-out posts it averages 2.55 bits per byte on an Apple M4, and 2.50 on a rerun on an NVIDIA RTX 3090 (its saved outputs): better than the LSTM, level with or slightly better than Kneser–Ney, and behind bz2.
- Taking pieces out: position is worth 0.68 bits per byte. Without residual connections the model collapses to predicting character frequencies.
- Its GPU utilization on the M4 is 25%: the matrix multiplies run at 68% of peak, and everything else is bound by memory.
Table of contents
Open Table of contents
The model
Here is the whole model, gpt.py less its imports and opening comment.
Each piece comes from an earlier post:
class Config:
def __init__(self, **kw):
self.V, self.T, self.D, self.L, self.N = 2048, 256, 192, 4, 4 # vocab, context, width, layers, heads
self.dropout, self.position, self.residual, self.prenorm = 0.1, "rope", True, True
self.__dict__.update(kw)
def rope(x): # x: (B, heads, T, H)
T, H = x.shape[-2], x.shape[-1]
ang = torch.arange(T, device=x.device)[:, None] * 10_000.0 ** (-torch.arange(0, H, 2, device=x.device) / H)
x1, x2 = x[..., 0::2], x[..., 1::2]
return torch.stack([x1 * ang.cos() - x2 * ang.sin(), x1 * ang.sin() + x2 * ang.cos()], -1).flatten(-2)
class RMSNorm(torch.nn.Module):
def __init__(self, D):
super().__init__(); self.g = torch.nn.Parameter(torch.ones(D))
def forward(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.g
class Block(torch.nn.Module):
def __init__(self, c):
super().__init__()
self.c, F_ = c, int(8 * c.D / 3) # SwiGLU hidden width: 8/3 D keeps the MLP at 8 D^2 weights
self.n1, self.n2 = RMSNorm(c.D), RMSNorm(c.D)
self.qkv, self.proj = torch.nn.Linear(c.D, 3 * c.D, bias=False), torch.nn.Linear(c.D, c.D, bias=False)
self.w_gate, self.w_up = torch.nn.Linear(c.D, F_, bias=False), torch.nn.Linear(c.D, F_, bias=False)
self.w_down = torch.nn.Linear(F_, c.D, bias=False)
self.drop = torch.nn.Dropout(c.dropout)
def attention(self, x):
B, T, D = x.shape
q, k, v = self.qkv(x).view(B, T, 3, self.c.N, D // self.c.N).permute(2, 0, 3, 1, 4)
if self.c.position == "rope":
q, k = rope(q), rope(k)
y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=self.c.dropout if self.training else 0)
return self.drop(self.proj(y.transpose(1, 2).reshape(B, T, D)))
def mlp(self, x):
return self.drop(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)))
def forward(self, x):
c = self.c
if not c.prenorm: # the 2017 layout: add, then normalize
x = self.n1(x + self.attention(x))
return self.n2(x + self.mlp(x))
if not c.residual: # no shortcut: each sublayer replaces the stream
x = self.attention(self.n1(x))
return self.mlp(self.n2(x))
x = x + self.attention(self.n1(x))
return x + self.mlp(self.n2(x))
class GPT(torch.nn.Module):
def __init__(self, c):
super().__init__()
self.c = c
self.embed = torch.nn.Embedding(c.V, c.D)
self.pos = torch.nn.Embedding(c.T, c.D) if c.position == "learned" else None
self.blocks = torch.nn.ModuleList([Block(c) for _ in range(c.L)])
self.norm = RMSNorm(c.D)
self.apply(self._init)
def _init(self, m):
if isinstance(m, (torch.nn.Linear, torch.nn.Embedding)):
torch.nn.init.normal_(m.weight, 0, 0.02)
def forward(self, idx):
x = self.embed(idx)
if self.pos is not None:
x = x + self.pos(torch.arange(idx.shape[1], device=idx.device))
for b in self.blocks:
x = b(x)
return self.norm(x) @ self.embed.weight.T # unembedding shares the embedding table
Reading it top to bottom, one position’s vector goes through:
- The embedding table, one row per token (see Embeddings and Contrastive Learning).
- blocks, each two sublayers, each read through an RMSNorm and
added back to the stream (see Dropout and Normalization and The Residual Stream):
- causal multi-head attention, heads of size , with
RoPE applied to the queries and keys (see The Transformer’s Three Attention Patterns and Where Am I? Position in Attention, from Sinusoids to RoPE). PyTorch’s
scaled_dot_product_attentioncomputes the softmax of with the causal mask, times , in one call, with dropout on the attention weights while training; - a SwiGLU MLP of hidden width (see The Residual Stream).
- causal multi-head attention, heads of size , with
RoPE applied to the queries and keys (see The Transformer’s Three Attention Patterns and Where Am I? Position in Attention, from Sinusoids to RoPE). PyTorch’s
- A final RMSNorm, then the unembedding, which here is the embedding table again, transposed. Sharing the table (weight tying, Press and Wolf, 2017) saves parameters and ties “what a token means going in” to “what predicts it coming out”.
This combination, RoPE, RMSNorm, SwiGLU and no bias terms, is the layout
of Llama (Touvron et al., 2023) and most open models since, at a
thousandth of the size. One simplification: every weight starts at a
standard deviation of 0.02, without GPT-2’s extra scaling of the matrices
that write into the stream. The three switches in Config, position,
residual and prenorm, are not part of a normal GPT. They’re there to
take pieces out later.
Counting the parameters
A block has weights in attention, one matrix each for queries, keys, values and the output, and in the MLP. With the embedding table that gives
plus the norm gains, each, which are too few to matter. Checked against the model, and against GPT-2:
from gpt import GPT, Config
# Where the parameters are, counted from the model and from the formula:
# per block 4 D^2 in attention (Q, K, V, output) and 8 D^2 in the SwiGLU MLP,
# plus the embedding table, which the unembedding shares.
c = Config(V=148) # the character model trained below
model = GPT(c)
groups = {"embedding (shared with unembedding)": [model.embed],
"attention, all blocks": [m for b in model.blocks for m in (b.qkv, b.proj)],
"MLP, all blocks": [m for b in model.blocks for m in (b.w_gate, b.w_up, b.w_down)],
"norm gains": [n for b in model.blocks for n in (b.n1, b.n2)] + [model.norm]}
total = sum(p.numel() for p in model.parameters())
for name, mods in groups.items():
n = sum(p.numel() for m in mods for p in m.parameters())
print(f"{name:37s} {n:>9,} {n / total:5.1%}")
print(f"{'total':37s} {total:>9,}")
formula = lambda L, D, V: 12 * L * D ** 2 + V * D
print(f"formula 12 L D^2 + V D: {formula(c.L, c.D, c.V):,} (L={c.L}, D={c.D}, V={c.V})")
print(f"with the 2,048 BPE tokens of What \"Generative\" Means instead: {formula(4, 192, 2048):,}, "
f"embedding {2048 * 192 / formula(4, 192, 2048):.0%} of it")
print(f"GPT-2 small (L=12, D=768, V=50,257): {formula(12, 768, 50257) / 1e6:.1f}M by the formula; "
f"124M reported, the rest mostly its 1,024 learned position vectors")
embedding (shared with unembedding) 28,416 1.6%
attention, all blocks 589,824 32.8%
MLP, all blocks 1,179,648 65.5%
norm gains 1,728 0.1%
total 1,799,616
formula 12 L D^2 + V D: 1,797,888 (L=4, D=192, V=148)
with the 2,048 BPE tokens of What "Generative" Means instead: 2,162,688, embedding 18% of it
GPT-2 small (L=12, D=768, V=50,257): 123.5M by the formula; 124M reported, the rest mostly its 1,024 learned position vectors
Where the parameters live. Two thirds are in the MLPs; the larger the vocabulary, the more of a small model is lookup table.
Data: the count aboveWith characters as tokens, is only 148, so the embedding is 1.6% of the model. With the 2,048 BPE tokens from What “Generative” Means it would be 18%, and in GPT-2 small, with 50,257 tokens, 31%. The larger the vocabulary, the more of a small model is lookup table. Looking up a row costs no arithmetic, but with tying the same table is also the unembedding, a real multiply, so the compute estimates below count every parameter: 1.8 million.
Training, with the choices made on a dev slice
The data is the frozen snapshot of the blog from Sequences Before Attention: 593,616 characters of training posts, and the two held-out posts, 57,381 bytes, that What “Generative” Means scored everything on. Those two posts are scored once, at the end. Every choice, the tokens, the dropout, the number of steps, the learning rate schedule, is made on a dev slice: the last 10% of the training text, with the model trained on the first 90%.
train.py holds the two splits, the two tokenizers (characters, and a
2,048-token byte-level BPE), the scoring and the training loop, and runs
one configuration from the command line. The loop:
def run(kind, seed, split, steps, checkpoints=(), B=32, lr=1e-3, warmup=200, schedule="constant",
device=DEVICE, on_checkpoint=None, **cfg):
V, encode, _ = tokenizer_for(kind, split["fit"])
fit, score = encode(split["fit"]).to(device), encode(split["score"]).to(device)
c = Config(V=V, **cfg)
torch.manual_seed(seed)
model = GPT(c).to(device)
decay = [p for p in model.parameters() if p.dim() == 2] # matrices decay; norm gains don't
rest = [p for p in model.parameters() if p.dim() < 2]
opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1}, {"params": rest, "weight_decay": 0.0}],
lr=lr, betas=(0.9, 0.95))
g = torch.Generator().manual_seed(seed)
curve = {}
for step in range(1, steps + 1):
frac = min(1.0, step / warmup)
if schedule == "cosine" and step > warmup:
frac = 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * (step - warmup) / (steps - warmup)))
for group in opt.param_groups:
group["lr"] = lr * frac
i = torch.randint(0, len(fit) - c.T - 1, (B,), generator=g).to(device)
idx = i[:, None] + torch.arange(c.T + 1, device=device)
x, y = fit[idx[:, :-1]], fit[idx[:, 1:]]
loss = F.cross_entropy(model(x).reshape(-1, V), y.reshape(-1))
opt.zero_grad(set_to_none=True); loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
if step in checkpoints or step == steps:
curve[step] = bits_per_byte(model, score, split["score"], c.T)
if on_checkpoint:
on_checkpoint(step, model)
passes = steps * B * c.T / len(fit)
return model, curve, passes
A few details that matter:
- Batches are 32 random windows of 256 tokens. Every position in a window is a training example, so one step trains on 8,192 predictions.
- The optimizer is AdamW at a peak rate of with , the usual choice for Transformers, weight decay 0.1 on the matrices (the tied embedding included) and not on the norm gains, a 200-step linear warmup and gradient clipping at 1 (see Initialization and Optimizers). The cosine schedule, when used, decays to 10% of the peak.
- Dropout is applied to the attention weights and to each sublayer’s output before it’s added to the stream (see Dropout and Normalization).
- Scoring slides a window of 256 by 128 and scores the second half of each window, so every character after the first 128 is predicted with at least 128 characters of context. The LSTM in What “Generative” Means carried its state through the whole text, so if anything this setup is harder on the GPT.
The first round of choices: characters or BPE tokens, and how much dropout.
One seed each, run by sweep.py, which calls train.py once per
configuration:
char d0.1 s0: 250: 2.741 / 500: 2.273 / 750: 2.158 / 1000: 2.076 / 1500: 2.080 / 2000: 2.072 / 2500: 2.105 / 3000: 2.141 (46 passes)
bpe d0.1 s0: 250: 2.415 / 500: 2.116 / 750: 2.204 / 1000: 2.396 / 1500: 2.791 / 2000: 3.066 / 2500: 3.218 / 3000: 3.333 (126 passes)
char d0.2 s0: 250: 2.919 / 500: 2.398 / 750: 2.230 / 1000: 2.147 / 1500: 2.093 / 2000: 2.053 / 2500: 2.040 / 3000: 2.034 (46 passes)
bpe d0.2 s0: 250: 2.461 / 500: 2.129 / 750: 2.111 / 1000: 2.177 / 1500: 2.394 / 2000: 2.547 / 2500: 2.693 / 3000: 2.789 (126 passes)
char d0.2 s0: 1000: 2.147 / 2000: 2.053 / 3000: 2.034 / 4000: 2.027 / 5000: 2.029 / 6000: 2.038 (92 passes)
char d0.3 s0: 1000: 2.254 / 2000: 2.116 / 3000: 2.088 / 4000: 2.030 / 5000: 2.039 / 6000: 2.015 (92 passes)
bpe d0.3 s0: 250: 2.522 / 500: 2.159 / 750: 2.106 / 1000: 2.114 / 1250: 2.169 / 1500: 2.201 / 2000: 2.263 (84 passes)
bpe d0.4 s0: 250: 2.603 / 500: 2.216 / 750: 2.128 / 1000: 2.117 / 1250: 2.119 / 1500: 2.122 / 2000: 2.164 (84 passes)
BPE tokens overfit fast: 2,048 tokens and 2.56 bytes per token mean only about 200,000 tokens of training text, so 2,000 steps is 84 passes, and the dev score is best around 750 steps and rises after. Characters give about 2.7 times as many training positions and keep improving to 4,000 steps. The best character runs, 2.02–2.03, beat the best BPE runs, 2.11, by more than the seed-to-seed spread measured below (0.01–0.02). So: characters, dropout 0.2, 4,000 steps. Dropout 0.3 at 6,000 steps was level (2.015 vs 2.027) and costs half as much again.
Taking pieces out
The second round, two seeds each, compares the chosen model with one change at a time, all at 4,000 steps on the dev slice. The first change is the learning rate schedule, still a choice to make; the rest are the ablations:
from sweep import train
# The second round, two seeds each, on the dev slice: the chosen model, then
# one change at a time (the learning-rate schedule, then the ablations).
BASE = dict(kind="char", steps=4000, dropout=0.2, checkpoints=[1000, 2000, 3000, 4000])
for seed in (0, 1):
train(f"baseline s{seed}", **BASE, seed=seed)
train(f"cosine schedule s{seed}", **BASE, seed=seed, schedule="cosine")
train(f"no position s{seed}", **BASE, seed=seed, position="none")
train(f"learned pos. s{seed}", **BASE, seed=seed, position="learned")
train(f"post-norm s{seed}", **BASE, seed=seed, prenorm=False)
train(f"no residual s{seed}", **BASE, seed=seed, residual=False)
baseline s0: 1000: 2.147 / 2000: 2.053 / 3000: 2.034 / 4000: 2.027 (61 passes)
cosine schedule s0: 1000: 2.141 / 2000: 2.016 / 3000: 2.001 / 4000: 1.990 (61 passes)
no position s0: 1000: 3.318 / 2000: 2.954 / 3000: 2.779 / 4000: 2.706 (61 passes)
learned pos. s0: 1000: 2.560 / 2000: 2.260 / 3000: 2.176 / 4000: 2.116 (61 passes)
post-norm s0: 1000: 2.292 / 2000: 2.127 / 3000: 2.095 / 4000: 2.088 (61 passes)
no residual s0: 1000: 5.061 / 2000: 5.022 / 3000: 5.049 / 4000: 5.048 (61 passes)
baseline s1: 1000: 2.140 / 2000: 2.030 / 3000: 2.023 / 4000: 2.006 (61 passes)
cosine schedule s1: 1000: 2.134 / 2000: 2.012 / 3000: 1.992 / 4000: 1.982 (61 passes)
no position s1: 1000: 3.374 / 2000: 2.918 / 3000: 2.790 / 4000: 2.688 (61 passes)
learned pos. s1: 1000: 2.563 / 2000: 2.281 / 3000: 2.169 / 4000: 2.111 (61 passes)
post-norm s1: 1000: 2.301 / 2000: 2.119 / 3000: 2.101 / 4000: 2.083 (61 passes)
no residual s1: 1000: 5.037 / 2000: 5.051 / 3000: 5.048 / 4000: 5.043 (61 passes)
What each piece is worth. Position matters most of the changes that still train; without residual connections the model does no better than counting characters.
Data: the output above- Cosine decay beats a constant rate on both seeds, 1.99 and 1.98 against 2.03 and 2.01. Lowering the learning rate at the end lets the weights settle into the bottom of the valley instead of bouncing across it (see Initialization and Optimizers). That settles the last choice. The rest of the ablations use the constant rate; each compares against the baseline trained the same way.
- Of the changes that still train, position matters most. With no position information the model scores 2.69–2.71, against the full model’s 2.01–2.03 on the same slice. It isn’t helpless, since the causal mask lets it count (see Where Am I? Position in Attention, from Sinusoids to RoPE), but it loses 0.68 bits per byte.
- Learned positions instead of RoPE: 2.11–2.12, 0.1 worse at 4,000 steps, and the gap is still closing: 0.41, 0.21, 0.14 and 0.09 at 1,000, 2,000, 3,000 and 4,000 steps (seed 0). There’s no extrapolation here, since training and scoring both use windows of 256, so what this measures is mostly speed: RoPE builds in that what matters is relative distance, and a learned table has to discover it, for each of 256 positions, from 600 KB of text. At this budget it hasn’t finished.
- Post-norm: 2.08–2.09, 0.06–0.08 worse. Normalizing after each addition puts a norm on the stream’s main path, the layout Dropout and Normalization describes as harder to optimize (Xiong et al., 2020). Two caveats: this variant got warmup and a learning rate tuned for pre-norm, and its first block reads the raw embedding, which starts tiny and is never normalized (the 2017 Transformer scaled it up by ). At 4 layers the gap may be as much about those as about the layout.
- No residual connections: 5.04–5.05, close to what a model that knows only how often each character occurs would score: the training text’s unigram entropy, printed below, is 5.015 bits per character. That needs a closer look.
What happens without the residual stream
A score at the unigram level suggests the model predicts the same distribution whatever the input. Checked directly, with the training loss and the gradient at each layer’s attention weights:
V, encode, _ = tokenizer_for("char", dev_split["fit"])
fit = encode(dev_split["fit"])
counts = torch.bincount(fit, minlength=V).float()
p = counts / counts.sum()
print(f"unigram entropy of the training text: {-(p[p > 0] * p[p > 0].log2()).sum():.3f} bits per character")
for residual in (True, False):
torch.manual_seed(0); g = torch.Generator().manual_seed(0)
m = GPT(Config(V=V, residual=residual, dropout=0.2))
opt = torch.optim.AdamW(m.parameters(), 1e-3, betas=(0.9, 0.95))
for step in range(1, 401):
for gr in opt.param_groups:
gr["lr"] = 1e-3 * min(1, step / 200)
i = torch.randint(0, len(fit) - 257, (16,), generator=g)
idx = i[:, None] + torch.arange(257)
x, y = fit[idx[:, :-1]], fit[idx[:, 1:]]
logits = m(x)
loss = F.cross_entropy(logits.reshape(-1, V), y.reshape(-1))
opt.zero_grad(); loss.backward()
if step in (1, 400):
norms = [m.blocks[l].qkv.weight.grad.norm().item() for l in range(4)]
print(f"residual={residual!s:5} step {step:3d}: train loss {loss.item() / math.log(2):.3f} bits; "
f"gradient norm at the attention weights of layers 1-4: "
+ " ".join(f"{n:.1e}" for n in norms))
torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
opt.step()
with torch.no_grad(): # does the output depend on the input at all?
m.eval()
a, b = m(x[:1])[0, -1].softmax(-1), m(x[1:2])[0, -1].softmax(-1)
print(f" output distributions for two different contexts differ by {0.5 * (a - b).abs().sum():.3f} (total variation)")
unigram entropy of the training text: 5.015 bits per character
residual=True step 1: train loss 7.214 bits; gradient norm at the attention weights of layers 1-4: 2.0e+00 1.7e+00 1.2e+00 9.3e-01
residual=True step 400: train loss 2.621 bits; gradient norm at the attention weights of layers 1-4: 5.6e-01 1.5e-01 1.5e-01 1.7e-01
output distributions for two different contexts differ by 0.861 (total variation)
residual=False step 1: train loss 7.246 bits; gradient norm at the attention weights of layers 1-4: 2.0e+00 1.4e+00 1.0e+00 7.5e-01
residual=False step 400: train loss 5.086 bits; gradient norm at the attention weights of layers 1-4: 1.6e-06 1.0e-06 3.3e-06 1.1e-04
output distributions for two different contexts differ by 0.000 (total variation)
The collapse, measured. Without residual connections the gradients fall by about four (layer 4) to six (layer 1) orders of magnitude by step 400, and the model’s prediction no longer depends on its input.
Data: the output aboveAt the first step the two models look alike: the gradients reaching the four layers are all of order 1, so this isn’t the vanishing gradient of Sequences Before Attention at initialization. By step 400 the model without residual connections has settled at 5.09 bits, near the unigram entropy, its gradients have fallen by about four (layer 4) to six (layer 1) orders of magnitude, and its prediction is identical for two different contexts. It has collapsed to the one output that ignores its input. The weights still move a little, since Adam divides each gradient by its own running size, but no step that ignores the input can lower the loss any further, and the tiny gradients that remain carry no useful direction. With residual connections, every sublayer’s output is added to a stream that still carries the input, so this state isn’t a trap. Is it just the settings? The same model for 1,000 steps with dropout off, and at lower learning rates (the same short loop on the CPU, batch 16, one seed, AdamW’s default weight decay):
lr 1e-03, dropout 0.0: train loss 4.966 bits after 1,000 steps; two contexts' predictions differ by 0.000
lr 3e-04, dropout 0.2: train loss 4.967 bits after 1,000 steps; two contexts' predictions differ by 0.000
lr 1e-04, dropout 0.2: train loss 4.726 bits after 1,000 steps; two contexts' predictions differ by 0.294
It collapses at and at . At it doesn’t collapse, and learns, slowly. So compare both models at that rate, for 4,000 steps, two seeds each (these runs, like the ablations, print scores on the dev slice):
from sweep import train
# The model without residual connections and the baseline, both at a tenth of
# the learning rate (1e-4), two seeds each, scored on the dev slice.
BASE = dict(kind="char", steps=4000, dropout=0.2, checkpoints=[1000, 2000, 3000, 4000])
for seed in (0, 1):
train(f"no residual, lr 1e-4 s{seed}", **BASE, seed=seed, residual=False, lr=0.0001)
for seed in (0, 1):
train(f"baseline, lr 1e-4 s{seed}", **BASE, seed=seed, lr=0.0001)
no residual, lr 1e-4 s0: 1000: 4.400 / 2000: 3.733 / 3000: 3.398 / 4000: 3.122 (61 passes)
no residual, lr 1e-4 s1: 1000: 4.373 / 2000: 3.660 / 3000: 3.328 / 4000: 3.041 (61 passes)
baseline, lr 1e-4 s0: 1000: 2.842 / 2000: 2.374 / 3000: 2.212 / 4000: 2.119 (61 passes)
baseline, lr 1e-4 s1: 1000: 2.837 / 2000: 2.337 / 3000: 2.169 / 4000: 2.102 (61 passes)
At the same learning rate, removing the residual connections costs about a bit per byte: 3.04–3.12 against 2.10–2.12, though the model without them is still improving fast. Comparing each model at its best rate gives the same answer, since the residual stream buys a second thing on top: it’s what lets the model train at ten times the rate, where it reaches 2.01–2.03 in the same number of steps. That is the practical meaning of the argument in The Residual Stream: the identity path doesn’t just help gradients reach early layers; it keeps the network out of states it can’t leave, so it can be trained fast.
The score, once
The chosen settings, characters, dropout 0.2, cosine decay over 4,000 steps, trained on all the training posts and scored on the two held-out posts. Three seeds, with the scores at earlier checkpoints printed too:
import sys
import torch
from train import run, final_split, tokenizer_for
# The one scored run: the settings chosen on the dev slice, trained on all
# the training posts, scored on the two held-out posts. Saves the model at
# a few checkpoints for sampling.
seed, schedule = int(sys.argv[1]), sys.argv[2]
vocab = tokenizer_for("char", final_split["fit"])[2]
def save(step, model):
if seed == 0:
state = {k: v.cpu() for k, v in model.state_dict().items()}
torch.save({"model": state, "config": dict(model.c.__dict__), "vocab": vocab}, f"gpt_{step}.pt")
_, curve, passes = run("char", seed, final_split, 4000, checkpoints=(250, 1000, 2000, 3000),
dropout=0.2, schedule=schedule, on_checkpoint=save)
print(f"seed {seed}, {schedule} schedule, held-out bits per byte at "
+ " / ".join(f"{s:,}: {v:.3f}" for s, v in curve.items()) + f" steps ({passes:.0f} passes)")
score_once.py runs it for seeds 0, 1 and 2:
seed 0, cosine schedule, held-out bits per byte at 250: 3.036 / 1,000: 2.511 / 2,000: 2.453 / 3,000: 2.475 / 4,000: 2.496 steps (55 passes)
seed 1, cosine schedule, held-out bits per byte at 250: 3.005 / 1,000: 2.607 / 2,000: 2.491 / 3,000: 2.562 / 4,000: 2.633 steps (55 passes)
seed 2, cosine schedule, held-out bits per byte at 250: 3.019 / 1,000: 2.492 / 2,000: 2.436 / 3,000: 2.570 / 4,000: 2.535 steps (55 passes)
At the fixed point chosen in advance, 4,000 steps, the three seeds score 2.50, 2.63 and 2.54 bits per byte, a mean of 2.55. The rerun on the RTX 3090 scores 2.50, 2.51 and 2.49, a mean of 2.50; a seed draws different random numbers on a different device, so these are three more samples, not a repeat. Against the ladder in What “Generative” Means:
- The LSTM, 2.66: beaten on all three seeds, by 0.03 to 0.16 on the M4 and 0.15 to 0.17 on the rerun. Its 3,000 steps of 64 windows of 128 characters trained on 25 million characters, close to this run’s 33 million, but with 380,000 parameters to this model’s 1.8 million, and without the dropout, weight decay and schedule tuned on a dev slice that this model got. So “the Transformer is better” here means “this Transformer, at this size, with this recipe”, not a matched comparison of architectures.
- Kneser–Ney, 2.57: on the M4, beaten on two seeds of three, and level on average; on the rerun, beaten on all three.
- bz2 given the training posts, 2.34: not beaten.
Two things in the checkpoints matter as much as the final numbers. On the held-out posts every seed was best at 2,000 or 3,000 steps, 2.44–2.49, then got worse, while on the dev slice the same settings kept improving to 4,000. On the M4 the seeds also drifted apart as they went, 0.06 apart at 2,000 steps and 0.14 at 4,000, but on the rerun they stayed within 0.04 of each other at every checkpoint from 2,000 on, so that drift is not a dependable property of the recipe. Chosen in hindsight, the 2,000-step checkpoint would beat Kneser–Ney on every seed. But that choice would have been made on the test set, which is what the dev slice exists to prevent, so the scored numbers are the pre-registered ones. (It’s also a checkpoint in the middle of a cosine schedule, at about 60% of the peak learning rate, not what a 2,000-step run would give.)
Why the held-out posts turned against more training is visible once they’re
scored separately. Every model on the ladder in What “Generative” Means, on each held-out post,
and the GPT (seed 0) also on 60,000 characters of the training text it
learned from. per_post.py does it with two helpers from other posts:
kn.py, the Kneser–Ney model from Sequences Before Attention, and lm.py, the LSTM from
What “Generative” Means. Run once as python lm.py char 0, lm.py trains that
post’s character-level LSTM again (seed 0) and saves it as char_lstm.pt. The LSTM row below came
from the original checkpoint; a retrained one can differ slightly. Each model’s cost for every
character is summed per post; for the GPT, at each checkpoint, and on
the training text too:
rows = {}
for step in (250, 1000, 2000, 3000, 4000):
ckpt = torch.load(f"gpt_{step}.pt", weights_only=False)
model = GPT(Config(**ckpt["config"])); model.load_state_dict(ckpt["model"]); model.eval()
bits = torch.cat([torch.zeros(1), bits_each(model, encode(val), model.c.T)])
seen_text = train[:60_000]
train_bpb = bits_each(model, encode(seen_text), model.c.T).sum().item() / nbytes(seen_text[1:])
rows[f"GPT, {step:,} steps"] = (per_post(bits), train_bpb)
if step == 4000:
cost = bits[unseen].sum().item()
print(f"the GPT at 4,000 steps pays {cost / unseen.sum():.1f} bits per unseen character, "
f"{cost / nbytes(val):.3f} bits per byte of the held-out text in all")
lstm = torch.load("char_lstm.pt", weights_only=False)
index = {u: i + 1 for i, u in enumerate(lstm["vocab"])}
m = LSTMLM(len(lstm["vocab"]) + 1); m.load_state_dict(lstm["model"]); m.eval()
data, bits, state = torch.tensor([index.get(c, 0) for c in val]), torch.zeros(len(val)), None
with torch.no_grad():
for s in range(0, len(data) - 1, 128):
x, y = data[s:s + 128], data[s + 1:s + 129]
logits, state = m(x[None][:, :len(y)], state)
bits[s + 1:s + 1 + len(y)] = F.cross_entropy(logits[0], y, reduction="none") / math.log(2)
rows["character LSTM, seed 0"] = (per_post(bits), None)
tables = build(6)
bits = torch.tensor([0.0] + [-math.log2(prob(tables, val[max(0, i - 6):i], val[i])) for i in range(1, len(val))])
rows["Kneser-Ney, 6 characters"] = (per_post(bits), None)
tb = train.encode()
rows["bz2 given the training posts"] = ({name: 8 * (len(bz2.compress(tb + val[a:b].encode(), 9)) - len(bz2.compress(tb, 9)))
/ nbytes(val[a:b]) for name, (a, b) in posts.items()}, None)
123 held-out characters never occur in training: BPE post 122, KV-cache post 1
the GPT at 4,000 steps pays 25.6 bits per unseen character, 0.055 bits per byte of the held-out text in all
bits per byte BPE post KV-cache post training text (seen)
GPT, 250 steps 3.212 2.654 2.698
GPT, 1,000 steps 2.775 1.935 1.803
GPT, 2,000 steps 2.747 1.815 1.566
GPT, 3,000 steps 2.801 1.763 1.398
GPT, 4,000 steps 2.842 1.744 1.327
character LSTM, seed 0 2.944 2.033
Kneser-Ney, 6 characters 2.826 2.003
bz2 given the training posts 2.480 2.038
The score, once, and why it hides two stories. On the post like its training text the GPT beats everything, bz2 included; on the post full of byte listings and code, bz2 wins.
Data: the scored runs and the per-post table aboveHover or focus a point to read its value.
The two held-out posts behave completely differently:
- The KV-cache post is prose about the same serving systems the training posts cover. On it the GPT keeps improving to the last checkpoint, 1.74 bits per byte, and beats every other model, bz2 included, by at least 0.26.
- The BPE post is full of things the training posts don’t contain: listings of byte tuples and merges, regular expressions, and 122 characters that never occur in training at all, mostly the visible-space symbol ␣ that tokenizer listings use, plus byte markers and an emoji. On it the GPT is best at 2,000 steps and then gets worse, and bz2 wins by 0.36.
The training-text column is the rest of the story: at 4,000 steps the model scores 1.33 bits per byte on text it has seen, against 1.74 and 2.84 on text it hasn’t, and the gap widens all the way. The model keeps fitting its training text more closely; that helps on a post like the training posts and hurts on one unlike them. The dev slice was like them, and more so than it looked: the 90% cut falls mid-sentence in one of the serving posts, so the dev slice begins by continuing a post the model trained on. A better dev split holds out whole posts, and different kinds of post. The lesson generalizes: a dev set predicts the test set only as well as it resembles it.
Why bz2 wins the BPE post
A file compressor from 1996 beats a Transformer by 0.36 bits per byte on one post. Three explanations, each measurable.
Characters it has never seen. The first line of the output above: 123 held-out characters, all but one in the BPE post, never occur in training. They all share the model’s “unseen” id, which it never saw as a target, so it pays 25.6 bits for each, 0.055 bits per byte of the held-out text in all, a third of seed 0’s gap to bz2. bz2 works on bytes, and a symbol it hasn’t seen costs it a few bytes of its own. A byte-level model, or a tokenizer that falls back to bytes (see Byte-Pair Encoding), wouldn’t have this problem.
Repetition. bz2 compresses the held-out posts together with the training posts, in blocks of 900 KB, and its Burrows–Wheeler transform sorts every position by what follows it, so text that repeats anything earlier, in training or in the same post, costs it little. The GPT sees only the last 256 characters. To measure how much text that is: mark each held-out character that sits inside a stretch of 32 (or 16) characters that appears word for word in the training posts (“copied”), or earlier in the held-out posts themselves, more than 256 characters back (“repeated”), and score the model on each kind. The masks are approximate: a stretch counts as repeated when its first occurrence is out of the window, even if a later copy sits inside it, and every character of a matched stretch is marked, including the first few, which no model could predict from the match.
def copied_mask(text, train, span):
covered = torch.zeros(len(text), dtype=torch.bool)
for i in range(len(text) - span + 1):
if text[i:i + span] in train:
covered[i:i + span] = True
return covered
def repeated_mask(text, span, T=256):
covered, first_seen = torch.zeros(len(text), dtype=torch.bool), {}
for i in range(len(text) - span + 1):
s = text[i:i + span]
if s in first_seen and first_seen[s] < i - T:
covered[i:i + span] = True
first_seen.setdefault(s, i)
return covered
repeats.py then scores the model at 4,000 steps on every character, as
train.py does, and averages the cost over each kind.
the model on the held-out posts: 2.522 bits per character
stretches of 32+ characters: copied 3.8% at 0.58 bits, repeated 1.6% at 5.15, the rest 94.7% at 2.56
stretches of 16+ characters: copied 15.4% at 1.07 bits, repeated 6.2% at 3.45, the rest 78.4% at 2.73
for comparison, the dev slice: 5.6% copied from the rest of the training text (32+)
What bz2 can see that the GPT can’t. Text repeated from beyond the model’s window is the most expensive text it reads.
Data: the output aboveThe model’s window shows in the numbers. On copied characters it pays 0.58 bits: it has learned those stretches, or enough of them. On repeated ones it pays 5.15, worse than on ordinary text, since they are the BPE post’s own unusual strings, its byte listings and code, which the model has no way to see again. Long repeats are rare: at 32 characters, 1.6% of the text, and coding them for free would save bits per character. But the answer depends on what counts as a repeat. At 16 characters, 6.2% of the text repeats something out of the model’s sight, at 3.45 bits each, worth bits, as much as the whole gap. bz2 exploits repeated contexts shorter still. So the model’s short window is a real part of why it loses: much of what bz2 knows about the BPE post, it learned from the BPE post itself.
Data. The third explanation is that the model hasn’t seen enough text. The same recipe on a quarter and on half of the training text, two seeds each:
import sys
from train import run, full_train, held_out
# The same model and schedule, trained on the first quarter or half of the
# training text, scored on the held-out posts: how much does more data help?
frac, seed = float(sys.argv[1]), int(sys.argv[2])
split = {"fit": full_train[:int(frac * len(full_train))], "score": held_out}
_, curve, passes = run("char", seed, split, 4000, checkpoints=(250, 500, 750, 1000, 2000, 3000), dropout=0.2, schedule="cosine")
print(f"{frac:.0%} of the training text ({len(split['fit']):,} characters), seed {seed}: held-out bits per byte at "
+ " / ".join(f"{s:,}: {v:.3f}" for s, v in curve.items()) + f" steps ({passes:.0f} passes)")
less_data.py runs it for each fraction and seed:
25% of the training text (148,404 characters), seed 0: held-out bits per byte at 250: 3.144 / 500: 2.979 / 750: 3.142 / 1,000: 3.308 / 2,000: 4.228 / 3,000: 4.878 / 4,000: 5.227 steps (221 passes)
25% of the training text (148,404 characters), seed 1: held-out bits per byte at 250: 3.142 / 500: 3.010 / 750: 3.212 / 1,000: 3.373 / 2,000: 4.178 / 3,000: 4.842 / 4,000: 5.216 steps (221 passes)
50% of the training text (296,808 characters), seed 0: held-out bits per byte at 250: 3.132 / 500: 2.869 / 750: 2.968 / 1,000: 2.916 / 2,000: 3.312 / 3,000: 3.535 / 4,000: 3.726 steps (110 passes)
50% of the training text (296,808 characters), seed 1: held-out bits per byte at 250: 3.168 / 500: 2.826 / 750: 2.915 / 1,000: 2.983 / 2,000: 3.080 / 3,000: 3.444 / 4,000: 3.608 steps (110 passes)
With less text the model overfits sooner, and every one of these runs is best early and then gets much worse, to 5.2 bits per byte for the quarter by 4,000 steps, 221 passes. So compare each at its best checkpoint, chosen in hindsight, the same way for all three:
- a quarter of the text: 2.98–3.01, at 500 steps;
- half: 2.83–2.87, at 500 steps;
- all of it: 2.44–2.49, at 2,000 steps (2,000 or 3,000 on the rerun).
More text helps, a lot, and the best point moves later, at roughly 14 to 28 passes over the data every time. (Muennighoff et al., 2023, found that up to about four passes over repeated data are nearly as good as new data, and that the value of more passes then decays toward zero; at 55 passes, the scored runs are far beyond that.) One caveat about the last doubling: the subsets are the first posts in alphabetical order, and the serving posts, vLLM, SGLang, TensorRT-LLM, come last, so the second half of the text holds nearly all of the training posts’ mentions of vLLM and of SGLang, the material closest to the held-out KV-cache post. So the step from half to all of it mixes more data with more relevant data. Either way, the model is starved: 600 KB is about what a large model sees in a fraction of a second of training, and bz2, which learns nothing in advance, needs no data to be good at repetition.
What the GPU is doing
Training cost, from Backpropagation, by Hand: about FLOPs per token, 2 for the forward pass and 4 for the backward, plus attention’s own work, the scores and the weighted sums, which don’t involve any parameters: per token for a context of , forward and backward (counting the full scores, as the PaLM paper’s convention does, though the mask makes half of them unnecessary). For this model, 10.8 million and 2.4 million per token, so attention’s own share is 18% at a context of 256. The run trains on 32.8 million tokens, about FLOPs in all.
Measured, and compared with the 2,878 GFLOP/s this GPU reached on a large
matrix multiply in How a GPU Runs Deep Learning. That peak, PEAK in
the script, is the M4’s; on another GPU it has to be set to that GPU’s own
measured peak, or the MFU means nothing (left unchanged on the RTX 3090,
the script prints 222%). mfu.py times 20
steps after 5 to warm up, waiting for the GPU (sync) before reading the
clock:
for kind in ("char", "bpe"):
V = tokenizer_for(kind, final_split["fit"])[0]
c = Config(V=V)
torch.manual_seed(0)
m = GPT(c).to(DEVICE)
opt = torch.optim.AdamW(m.parameters(), 1e-3)
B = 32
x = torch.randint(0, V, (B, c.T + 1), device=DEVICE)
n_params = sum(p.numel() for p in m.parameters())
for i in range(25):
if i == 5:
sync(); t0 = time.perf_counter()
loss = F.cross_entropy(m(x[:, :-1]).reshape(-1, V), x[:, 1:].reshape(-1))
opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
sync()
step = (time.perf_counter() - t0) / 20
tokens = B * c.T
flops = tokens * (6 * n_params + 12 * c.L * c.T * c.D)
print(f"{kind:4s}: {n_params:,} parameters, {step * 1e3:.0f} ms per step of {tokens:,} tokens, "
f"{tokens / step:,.0f} tokens/s, {flops / step / 1e9:.0f} GFLOP/s, MFU {flops / step / PEAK:.1%}")
char: 1,799,616 parameters, 149 ms per step of 8,192 tokens, 54,915 tokens/s, 723 GFLOP/s, MFU 25.1%
bpe : 2,164,416 parameters, 161 ms per step of 8,192 tokens, 50,726 tokens/s, 778 GFLOP/s, MFU 27.0%
About 25% model FLOPs utilization (MFU): of every four FLOPs the GPU could do, the model’s arithmetic uses one. A step of 8,192 tokens takes about 150 ms, so the 4,000 steps take 10 minutes, plus a few more for scoring the checkpoints. The BPE model, with a larger unembedding matrix, does slightly better.
Where do the other 75% go? It’s tempting to blame small matrices, but
every multiply here has 8,192 rows, one per token in the batch, which is
plenty of work per call. Timing the parts of a step separately, each
with timed, which runs a function 3 times to warm up and then averages
20 runs:
c = Config(V=148)
M, D, F_ = 32 * c.T, c.D, 512
shapes = [(D, 3 * D), (D, D), (D, F_), (D, F_), (F_, D)] * c.L + [(D, c.V)] # per layer, then unembedding
xs = [torch.randn(M, k, device=dev, requires_grad=True) for k, _ in shapes]
ws = [torch.randn(k, n, device=dev, requires_grad=True) for k, n in shapes]
def matmuls():
sum((x @ w).sum() for x, w in zip(xs, ws)).backward()
flops = sum(6 * M * k * n for k, n in shapes) # 2 forward + 4 backward per multiply-add
t = timed(matmuls)
print(f"the {len(shapes)} matrix multiplies alone, forward and backward: {t:.0f} ms, {flops / t / 1e6:.0f} GFLOP/s")
q, k, v = (torch.randn(32, c.N, c.T, D // c.N, device=dev, requires_grad=True) for _ in range(3))
for p in (0.0, 0.2):
def attn():
F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=p).sum().backward()
print(f"attention's fused call, {c.L} layers, dropout {p}: {c.L * timed(attn):.0f} ms")
the 21 matrix multiplies alone, forward and backward: 45 ms, 1960 GFLOP/s
attention's fused call, 4 layers, dropout 0.0: 40 ms
attention's fused call, 4 layers, dropout 0.2: 53 ms
AdamW's update of all parameters: 1.5 ms
Where a step’s time goes. The parameters’ multiplies take less than a third of it; the rest is moving memory.
Data: the timings aboveThe 21 matrix multiplies, all the model’s parameters at work, take 45 ms of the 150, at 68% of the peak. Attention’s own call takes 53 ms, a third of the step, for less than a fifth of the FLOPs: it computes and stores a score matrix per head and sequence, then the softmax and, in training, a dropout mask on it, which is mostly moving memory (FlashAttention is built to avoid exactly that). The optimizer takes under 2 ms. The remaining 50 ms or so are the operations that do little arithmetic per byte: the norms, RoPE, the SiLU and the gate, the dropout on each sublayer, the loss, all bound by memory (see Dropout and Normalization and How a GPU Runs Deep Learning). Large training runs on data-center accelerators report MFU around 40%: 46% for PaLM on TPUs, 38–43% for Llama 3 405B on H100s. They get there with fused kernels for exactly these memory-bound pieces, and with matrices large enough that the multiplies dominate.
What it writes, and what it copied
Samples from the scored model (seed 0) after 250, 1,000 and 4,000 steps, all from the same prompt, at temperature 0.8 (see What “Generative” Means), with the longest stretch of each that appears word for word in the training posts:
def longest_copied(text):
best, i = "", 0
for i in range(len(text)):
j = i + len(best) + 1
while j <= len(text) and text[i:j] in train:
best, j = text[i:j], j + 1
return best
@torch.no_grad()
def sample(path, prompt, n=300, temperature=0.8, seed=0, greedy=False):
ckpt = torch.load(path, weights_only=False)
c = Config(**ckpt["config"])
model = GPT(c); model.load_state_dict(ckpt["model"]); model.eval()
vocab = ckpt["vocab"]
index = {ch: i + 1 for i, ch in enumerate(vocab)}
g = torch.Generator().manual_seed(seed)
ids = [index.get(ch, 0) for ch in prompt]
for _ in range(n):
logits = model(torch.tensor([ids[-c.T:]]))[0, -1] / temperature
ids.append(logits.argmax().item() if greedy else
torch.multinomial(torch.softmax(logits, -1), 1, generator=g).item())
return "".join(vocab[i - 1] if i > 0 else "?" for i in ids)
--- after 250 steps (temperature 0.8)
The KV cache an the rownumes the Cone
the this
***Kuad models is small in the is the methat comput, a abl.
- **The wult it's approduction the rough with recodes log-own is so
matches the model's every over *attention the loss
blocks is every the the revery in that training, un every eform in model are math mo
[longest stretch copied from the training posts: 14 characters: "es the model's"]
--- after 1,000 steps (temperature 0.8)
The KV cache context of units to Context at it
sequence one scheduler be has is all attention 0.1. The blog loss context
process the one He in the whole recosystem to computed to the $W_V$, the
with the other attention catch from the three with the one number of $\matrix 1$.
The gradient. The model can move mo
[longest stretch copied from the training posts: 17 characters: ' from the three w']
--- after 4,000 steps (temperature 0.8)
The KV cache post
1's own exactly it can always a KV cache per layer, because it can be what
predicts above a fixed-size process that keeps the fixed prompt per log-odds
per-token context sequences and doesn't change, and dividing this post's
change in the next term training, use the serving posts are predictin
[longest stretch copied from the training posts: 19 characters: ' the serving posts ']
held-out text from character 10,000: longest stretch also in the training posts 27 characters: 's |\n| --- | --- | --- |\n| `'
held-out text from character 25,000: longest stretch also in the training posts 17 characters: '.\n\n### Where the '
held-out text from character 40,000: longest stretch also in the training posts 70 characters: ' = \\operatorname{softmax}\\!\\left(\\frac{QK^\\top}{\\sqrt{d_k}} + M\\right)'
ten more samples at 4,000 steps: longest copied stretch 15 to 38 characters
greedy, from a code prompt:
```python
import torch
from router_data import split, featurize
from overfit import sigmoid, log_loss, train
# AdamW that starts the same per-step, the same per-sequence sequence length
# structured and context problem tha
[longest stretch copied: 89 characters: 'from router_data import split, featurize\nfrom overfit import sigmoid, log_loss, train\n\n# ']
At 250 steps the model has spelling-level structure: common short words, markdown bullets and bold, a line length about right. At 1,000 it has phrases (“the whole”, “attention”, "") and LaTeX. At 4,000 it writes grammatical fragments in the blog’s register, about KV caches, per-layer context and serving, that don’t add up to a sentence with a meaning. That is roughly what 2.5 bits per character sounds like, sharpened a little by the temperature.
Has it memorized the blog? Its prose, no. The longest stretch of these samples found word for word in the training posts is 14 to 19 characters, phrases like “the serving posts”, and ten more samples give 15 to 38. Sampling draws different random numbers on another device, so the rerun’s samples differ: 14, 15 and 64 characters, the 64 an import line from a code block, and 15 to 54 for ten more. For comparison, three pieces of the held-out posts, text a person wrote, share 17 to 70 characters with the training posts, since the blog repeats its own headings, table syntax and formulas.
Its boilerplate, yes. Greedy decoding from the start of a code block reproduces 89 characters verbatim (48 on the rerun): import lines that recur across the training posts’ code blocks, the model’s most reliable text. And the model’s scores say it has fitted the training text far more closely than its samples show: 1.33 bits per byte on text it trained on, against 2.5 on new text. What it has memorized is what the blog repeats; everything else it has learned only as statistics.
Where this leads
- Decoding turns the model’s distributions into text: greedy decoding, temperature, top-k and top-p, beam search, and why they differ.
- How LLM Architecture Evolved, 2017–2026 covers scaling laws, the question this post ran into: for a fixed budget, how large a model, on how much data?
- The KV Cache removes the waste in the sampling script above, which recomputes all 256 positions for every new character.
The assistant has its first real language model: a 1.8-million-parameter GPT trained on the notebook itself, scoring 2.50–2.63 bits per byte (mean 2.55) on two posts it has never seen, and 2.49–2.51 on a rerun on another GPU. It beats bz2 on the post that resembles its training text and loses on the one that doesn’t, and it writes text that looks like the notebook and means nothing. The next parts give it better decoding, then more data and a bigger model, before it’s asked to be useful.
Takeaways
- A GPT is a short list of earlier pieces: an embedding table, blocks of normalized attention and MLP added to a residual stream, a final norm and a (tied) unembedding. With RoPE, RMSNorm and SwiGLU it’s Llama’s layout, in about 80 lines.
- Its parameters are plus ; the formula matches the model to within the norm gains, and GPT-2 small to within its position table. Training costs about FLOPs per token, plus for attention’s own work.
- Choose on a dev slice, score the test set once, and report what was fixed in advance. Here the held-out posts preferred an earlier checkpoint than the dev slice did, and the honest score is the pre-registered one: 2.55 bits per byte on average on the M4 and 2.50 on the rerun, better than the LSTM (2.66), level with or slightly better than Kneser–Ney (2.57), behind bz2 (2.34).
- Break an aggregate down before explaining it. Scored per post, the GPT beats bz2 by 0.3 on the post like its training text and loses by 0.36 on the one unlike it, a loss made of never-seen characters, repeats beyond its window, and too little data.
- Ablations, at a fixed budget on the dev slice: no position information costs 0.68 bits per byte, learned positions 0.1 (and closing), post-norm 0.06–0.08. Without residual connections the model collapses to the unigram distribution at a normal learning rate, and at one ten times smaller it still trails by about a bit.
- MFU on the M4 is 25%. The matrix multiplies run at 68% of the peak; the rest of the step is attention’s score matrix and memory-bound elementwise work, the parts fused kernels exist to speed up.
References
Start with Karpathy’s nanoGPT and its video walkthrough, which this model closely follows. The Llama paper describes the RoPE, RMSNorm and SwiGLU layout used here, and the PaLM paper defines MFU.
- Radford, Narasimhan, Salimans and Sutskever, Improving Language Understanding by Generative Pre-Training (OpenAI, 2018): GPT, a decoder-only Transformer trained on next-token prediction.
- Radford et al., Language Models are Unsupervised Multitask Learners (OpenAI, 2019): GPT-2 and pre-norm. The paper lists its smallest model as 117M parameters; the released checkpoint counted above has 124M.
- Press and Wolf, Using the Output Embedding to Improve Language Models (EACL, 2017): weight tying.
- Shazeer, GLU Variants Improve Transformer (2020): SwiGLU.
- Zhang and Sennrich, Root Mean Square Layer Normalization (NeurIPS, 2019): RMSNorm.
- Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021): RoPE.
- Muennighoff et al., Scaling Data-Constrained Language Models (NeurIPS, 2023): what repeated passes over the same data are worth.
- Grattafiori et al. (Meta AI), The Llama 3 Herd of Models (2024): 38–43% MFU training the 405B model.
- Touvron et al. (Meta AI), LLaMA: Open and Efficient Foundation Language Models (2023): the RoPE, RMSNorm and SwiGLU layout used here.
- Xiong et al., On Layer Normalization in the Transformer Architecture (ICML, 2020): why post-norm needs warmup and pre-norm doesn’t.
- Kaplan et al., Scaling Laws for Neural Language Models (2020): the FLOPs per token approximation.
- Chowdhery et al., PaLM: Scaling Language Modeling with Pathways (2022): model FLOPs utilization, including attention’s term.
- Karpathy, nanoGPT: a compact GPT training codebase in the same spirit.