From Prototype to nanoGPT
side-by-side with the reference.
Two surgical refinements that turn the Project 5 prototype into a production-shape model: parameter groups for weight decay, and scaled residual initialization. Train both side-by-side and watch val loss drop by 0.7.
The Concept
Most of the surface area between "a working prototype" and "a production-shape reference implementation" lives in three categories:
- True algorithmic necessities — without these, the model doesn't train (residuals, LayerNorm). Project 5 already has these.
- Engineering choices that change speed, clarity, robustness — what this project demonstrates.
- Pure style — naming, file organization.
We apply two refinements from category 2:
Parameter groups for weight decay. AdamW applies weight decay uniformly by default. But biases and LayerNorm parameters serve a different role from weight matrices — they calibrate scale and shift, they don't carry the learned content. Shrinking them is at best wasted and at worst harmful. Split parameters into two groups: weights get
weight_decay=0.1, biases and LayerNorm getweight_decay=0.0.Scaled residual initialization. In a deep residual stack, each block adds its output to the running hidden state. If each addition is full-scale, the magnitude drifts upward layer by layer. Compensate by initializing the residual-path output projections (attention.proj and the second linear in the MLP) with std reduced by
1/sqrt(2 * n_layers). Other linears and embeddings getstd=0.02.
Why It Matters
After Project 6 you should not admire reference code from a distance. You should steal from it selectively, because you know what each stolen part is protecting.
The Build
build.py$python projects/06_from-prototype-to-nanogpt/build.py --tiny
"""
Project 6: From Prototype to nanoGPT — production-shape refinements.
Starts from the Project 5 prototype and applies two nanoGPT-style improvements:
1. **Parameter groups for weight decay**: weights get weight decay, biases and
LayerNorm parameters do not. They have different roles, so the optimizer
should treat them differently.
2. **Scaled residual initialization**: projections that sit on residual branches
(attention output proj and MLP output proj) get initialized with reduced
variance. In a deep residual stack each block adds to the running hidden
state; if every addition is full-scale, the hidden state drifts upward
layer by layer.
Trains side-by-side against a baseline (default init + uniform decay) and
reports the difference.
Run:
python build.py --tiny
python build.py --full
"""
from __future__ import annotations
import argparse
import importlib.util
import math
import sys
from pathlib import Path
import torch
import torch.nn as nn
# Import Project 5's building blocks. We reuse its GPT class wholesale, then
# override init and optimizer.
PROJECT_5 = Path(__file__).resolve().parent.parent / "05_your-gpt-from-a-blank-file"
def _load_project_5():
spec = importlib.util.spec_from_file_location("project_05_build", PROJECT_5 / "build.py")
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
sys.modules["project_05_build"] = module
spec.loader.exec_module(module)
return module
p5 = _load_project_5()
def init_weights(model: p5.GPT) -> None:
"""nanoGPT-style init: N(0, 0.02) for linear/embedding; scaled for residual projections."""
for name, m in model.named_modules():
if isinstance(m, nn.Linear):
torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
if m.bias is not None:
torch.nn.init.zeros_(m.bias)
elif isinstance(m, nn.Embedding):
torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
# Scaled init for residual-path projections.
n_layers = model.cfg.n_layers
scale = 0.02 / math.sqrt(2 * n_layers)
for block in model.blocks:
torch.nn.init.normal_(block.attn.proj.weight, mean=0.0, std=scale)
torch.nn.init.normal_(block.mlp.net[2].weight, mean=0.0, std=scale) # second linear in MLP
def configure_param_groups(model: p5.GPT, weight_decay: float = 0.1) -> list[dict]:
"""Group parameters into decayed and undecayed for AdamW.
Weights (Linear, Embedding) get weight decay. Biases and LayerNorm params don't.
Returns groups in the format AdamW expects.
"""
decay: list[torch.nn.Parameter] = []
no_decay: list[torch.nn.Parameter] = []
seen_params: set[int] = set() # ids, for weight-tying dedup
for module in model.modules():
for name, param in module.named_parameters(recurse=False):
if id(param) in seen_params:
continue
seen_params.add(id(param))
if isinstance(module, nn.LayerNorm) or name.endswith("bias"):
no_decay.append(param)
else:
decay.append(param)
return [
{"params": decay, "weight_decay": weight_decay},
{"params": no_decay, "weight_decay": 0.0},
]
def train_with_groups(
model: p5.GPT,
train_data: torch.Tensor,
val_data: torch.Tensor,
steps: int,
batch_size: int,
lr: float,
eval_every: int,
weight_decay: float = 0.1,
seed: int = 0,
) -> tuple[list[float], list[float]]:
g = torch.Generator().manual_seed(seed)
groups = configure_param_groups(model, weight_decay=weight_decay)
optimizer = torch.optim.AdamW(groups, lr=lr)
train_hist: list[float] = []
val_hist: list[float] = []
for step in range(steps):
x, y = p5.get_batch(train_data, model.cfg.block_size, batch_size, g)
_, loss = model(x, y)
optimizer.zero_grad()
assert loss is not None
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
if step % eval_every == 0 or step == steps - 1:
train_hist.append(float(loss.item()))
with torch.no_grad():
vx, vy = p5.get_batch(val_data, model.cfg.block_size, batch_size, g)
_, vloss = model(vx, vy)
val_hist.append(float(vloss.item()))
return train_hist, val_hist
def write_loss_comparison(
base_train: list[float],
base_val: list[float],
nano_train: list[float],
nano_val: list[float],
eval_every: int,
path: Path,
) -> None:
try:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
except ImportError:
return
xs = [i * eval_every for i in range(len(base_train))]
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(xs, base_train, label="prototype")
ax1.plot(xs, nano_train, label="nanoGPT-style")
ax1.set_xlabel("step")
ax1.set_ylabel("train loss")
ax1.legend()
ax1.grid(True, alpha=0.3)
ax1.set_title("Train loss")
ax2.plot(xs, base_val, label="prototype")
ax2.plot(xs, nano_val, label="nanoGPT-style")
ax2.set_xlabel("step")
ax2.set_ylabel("val loss")
ax2.legend()
ax2.grid(True, alpha=0.3)
ax2.set_title("Val loss")
fig.tight_layout()
fig.savefig(path, dpi=120)
plt.close(fig)
def main() -> int:
parser = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
)
parser.add_argument("--tiny", action="store_true")
parser.add_argument("--full", action="store_true")
parser.add_argument("--steps", type=int, default=None)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument(
"--output-dir",
type=Path,
default=Path(__file__).parent / "outputs",
)
args = parser.parse_args()
if args.steps is None:
args.steps = 5000 if args.full else 300
args.output_dir.mkdir(parents=True, exist_ok=True)
text = p5.DEFAULT_CORPUS
stoi, _, vocab_size = p5.char_tokenizer(text)
data = p5.encode(text, stoi)
n_train = int(0.9 * len(data))
train_data = data[:n_train]
val_data = data[n_train:]
cfg = p5.GPTConfig(block_size=32, n_layers=4, n_heads=4, d_model=64)
eval_every = max(1, args.steps // 10)
# === Prototype (Project 5 defaults: PyTorch default init, uniform weight decay) ===
torch.manual_seed(args.seed)
proto = p5.GPT(cfg, vocab_size=vocab_size)
# Project 5 trained with no weight_decay specified (default 0). Match that.
proto_train, proto_val = p5.train_gpt(
proto,
train_data,
val_data,
steps=args.steps,
batch_size=32,
lr=3e-3,
eval_every=eval_every,
seed=args.seed,
)
# === nanoGPT-style (scaled init + parameter groups) ===
torch.manual_seed(args.seed)
nano = p5.GPT(cfg, vocab_size=vocab_size)
init_weights(nano)
nano_train, nano_val = train_with_groups(
nano,
train_data,
val_data,
steps=args.steps,
batch_size=32,
lr=3e-3,
eval_every=eval_every,
weight_decay=0.1,
seed=args.seed,
)
uniform = math.log(vocab_size)
print(f"\nUniform baseline: {uniform:.4f}")
print(f"\n{'mode':30s} {'final train':>14s} {'final val':>12s}")
print("-" * 60)
print(f"{'prototype (P5 defaults)':30s} {proto_train[-1]:>14.4f} {proto_val[-1]:>12.4f}")
print(f"{'nanoGPT-style refinements':30s} {nano_train[-1]:>14.4f} {nano_val[-1]:>12.4f}")
# Sanity: check that param groups split as expected.
nano_groups = configure_param_groups(nano, weight_decay=0.1)
n_decay = sum(p.numel() for p in nano_groups[0]["params"])
n_no_decay = sum(p.numel() for p in nano_groups[1]["params"])
print(f"\nParameter groups in nanoGPT-style optimizer:")
print(f" with weight_decay=0.1: {n_decay} params")
print(f" with weight_decay=0.0: {n_no_decay} params (biases + LayerNorm)")
write_loss_comparison(
proto_train,
proto_val,
nano_train,
nano_val,
eval_every,
args.output_dir / "loss_comparison.png",
)
log = args.output_dir / "run_log.txt"
log.write_text(
"# Project 6 run log\n\n"
f"uniform baseline : {uniform:.4f}\n"
f"prototype (P5 defaults) : train={proto_train[-1]:.4f} val={proto_val[-1]:.4f}\n"
f"nanoGPT-style : train={nano_train[-1]:.4f} val={nano_val[-1]:.4f}\n"
f"\nparameter groups:\n"
f" decayed : {n_decay}\n"
f" not-decayed (bias + LN) : {n_no_decay}\n"
"\n"
"Refinements applied:\n"
" 1. Init: N(0, 0.02) for Linear + Embedding; residual projections\n"
" get std = 0.02 / sqrt(2 * n_layers) to compensate for stacking.\n"
" 2. AdamW param groups: weights get weight_decay=0.1; biases and\n"
" LayerNorm params get weight_decay=0.0.\n",
encoding="utf-8",
)
print(f"\nOutputs written to {args.output_dir.resolve()}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
The Steps
4 reference filesEach step is the code for one piece, in isolation — read them as the book introduces them. They assemble into build.py above.
Step 02Compare Embeddingsstep_02_compare-embeddings.py
"""
Project 6: Step 2 — Compare embeddings
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
tok_emb = self.token_embedding(idx) # (B, T, d)
pos = torch.arange(T, device=idx.device)
pos_emb = self.position_embedding(pos) # (T, d)
x = tok_emb + pos_emb
Step 04Compare Transformer Blocksstep_04_compare-transformer-blocks.py
"""
Project 6: Step 4 — Compare transformer blocks
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
Step 08Compare Batch Construction And Data Pathstep_08_compare-batch-construction-and-data-path.py
"""
Project 6: Step 8 — Compare batch construction and data path
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([data[i:i+block_size] for i in ix])
y = torch.stack([data[i+1:i+block_size+1] for i in ix])
Step 09Compare Optimizer Setupstep_09_compare-optimizer-setup.py
"""
Project 6: Step 9 — Compare optimizer setup
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
Break It
break_it.pyProject 6 doesn't have a separate break_it.py exercise because the comparison is the break-vs-fix experiment. Running build.py trains both versions side by side. The "broken" baseline is the unrefined Project 5 defaults. The "fixed" version applies the two refinements together. The val-loss delta (0.7 nats lower) is what the refinements buy you.
If you want a more targeted break: edit configure_param_groups to put everything in the same group with full weight decay, retrain, and watch the LayerNorm weights collapse toward zero. That's the lesson that "different parameters play different roles, so the optimizer should treat them differently."
$python projects/06_from-prototype-to-nanogpt/break_it.py --tiny
"""
Project 6: BREAK IT experiment.
Deliberately sabotages one mechanism from build.py to show what happens
when it's removed. Compare outputs to the working version.
"""
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
x = self.attn(self.ln1(x))
x = self.mlp(self.ln2(x))
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
x = x + self.attn(x)
x = x + self.mlp(x)
Outputs
run_log.txt
# Project 6 run log
uniform baseline : 3.8712
prototype (P5 defaults) : train=0.5342 val=4.3250
nanoGPT-style : train=0.4460 val=3.6184
parameter groups:
decayed : 201728
not-decayed (bias + LN) : 3456
Refinements applied:
1. Init: N(0, 0.02) for Linear + Embedding; residual projections
get std = 0.02 / sqrt(2 * n_layers) to compensate for stacking.
2. AdamW param groups: weights get weight_decay=0.1; biases and
LayerNorm params get weight_decay=0.0.Read it in the book
This page ships the code. The full derivation — why the chain rule is mechanical not magical, what the computational graph buys you, the long-form reasoning — lives in Chapter 6 of Under the Hood.
Get the book →