The Details That Matter
norms, activations, positional encodings.
RMSNorm instead of LayerNorm. SwiGLU instead of GELU. The two architectural swaps that appear in Llama, Mistral, and most modern open-weight LLMs. Train all four combinations side-by-side and look at the deltas.
The Concept
- RMSNorm. LayerNorm computes mean and variance, then normalizes. RMSNorm drops the mean step and only divides by root-mean-square:
x / sqrt(mean(x²) + eps). Same effect on signal magnitude; about half the work. - SwiGLU. Standard MLP:
x → Linear → GELU → Linear. SwiGLU adds a gating term:(silu(x @ W_gate) * (x @ W_up)) @ W_down. The gating multiplies the projected representation by a learned mask — features get scaled by their own activation, which makes the MLP more expressive without changing parameter count much (we set hidden_mult=8/3 to keep counts roughly equal to the GELU MLP).
Why It Matters
These are not algorithmic breakthroughs. They are incremental refinements that win on the margins. Reading reference implementations of modern models without knowing what's underneath these names is reading code with two unexplained imports.
The Build
build.py$python projects/07_the-details-that-matter/build.py --tiny
"""
Project 7: The Details That Matter — RMSNorm and SwiGLU as alternatives.
Two architectural details that show up in modern LLMs (Llama, Mistral, etc.)
as replacements for the LayerNorm + GELU from the earlier transformer recipe:
1. **RMSNorm** — drops LayerNorm's mean-centering. Just normalize by root-mean-square.
Same regularization effect at ~half the FLOPs.
2. **SwiGLU** — replaces the MLP's `Linear -> GELU -> Linear` with a gated form:
`(silu(x @ W1) * (x @ W2)) @ W3`. The gating term learns which features to
pass through. Slightly more parameters; consistently better in modern LLMs.
This project trains a tiny GPT in four configurations and measures the
difference: {LayerNorm, RMSNorm} x {GELU MLP, SwiGLU MLP}.
Run:
python build.py --tiny # ~30s on CPU
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 torch.nn.functional as F
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()
class RMSNorm(nn.Module):
"""Root-mean-square layer norm. No mean centering; only divide by RMS."""
def __init__(self, d: int, eps: float = 1e-6) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(d))
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
return self.weight * (x / rms)
class SwiGLU(nn.Module):
"""Gated MLP: out = (silu(x @ W1) * (x @ W2)) @ W3.
Standard GELU MLP: d -> 4d -> d (= 8 d^2 params)
SwiGLU: d -> 4d (gate) + d -> 4d (proj), then 4d -> d
= 4d^2 + 4d^2 + 4d^2 = 12d^2 params (with full 4d hidden)
To keep parameter count roughly equal to a GELU MLP, we use 8/3 * d for the
hidden dim, which gives roughly 8 d^2 params total. We use floor to keep
things divisible.
"""
def __init__(self, d: int, hidden_mult: float = 8.0 / 3.0) -> None:
super().__init__()
d_hidden = int(d * hidden_mult)
self.gate_proj = nn.Linear(d, d_hidden, bias=False)
self.up_proj = nn.Linear(d, d_hidden, bias=False)
self.down_proj = nn.Linear(d_hidden, d, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
class ModernBlock(nn.Module):
"""Transformer block with configurable norm and MLP."""
def __init__(self, cfg: p5.GPTConfig, norm_type: str = "ln", mlp_type: str = "gelu") -> None:
super().__init__()
Norm = RMSNorm if norm_type == "rms" else nn.LayerNorm
self.ln1 = Norm(cfg.d_model)
self.attn = p5.CausalSelfAttention(cfg)
self.ln2 = Norm(cfg.d_model)
self.mlp = SwiGLU(cfg.d_model) if mlp_type == "swiglu" else p5.MLP(cfg)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
class ModernGPT(nn.Module):
"""Tiny GPT with configurable normalization and MLP variants."""
def __init__(
self,
cfg: p5.GPTConfig,
vocab_size: int,
norm_type: str = "ln",
mlp_type: str = "gelu",
) -> None:
super().__init__()
self.cfg = cfg
self.vocab_size = vocab_size
Norm = RMSNorm if norm_type == "rms" else nn.LayerNorm
self.token_embedding = nn.Embedding(vocab_size, cfg.d_model)
self.position_embedding = nn.Embedding(cfg.block_size, cfg.d_model)
self.blocks = nn.ModuleList(
[ModernBlock(cfg, norm_type=norm_type, mlp_type=mlp_type) for _ in range(cfg.n_layers)]
)
self.ln_f = Norm(cfg.d_model)
self.lm_head = nn.Linear(cfg.d_model, vocab_size, bias=False)
self.lm_head.weight = self.token_embedding.weight
def forward(
self, idx: torch.Tensor, targets: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor | None]:
B, T = idx.shape
pos = torch.arange(T, device=idx.device)
x = self.token_embedding(idx) + self.position_embedding(pos)
for block in self.blocks:
x = block(x)
x = self.ln_f(x)
logits = self.lm_head(x)
if targets is None:
return logits, None
loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1))
return logits, loss
def train_variant(
norm_type: str,
mlp_type: str,
vocab_size: int,
train_data: torch.Tensor,
val_data: torch.Tensor,
steps: int,
seed: int = 0,
) -> tuple[ModernGPT, float, float]:
torch.manual_seed(seed)
cfg = p5.GPTConfig(block_size=32, n_layers=2, n_heads=4, d_model=64)
model = ModernGPT(cfg, vocab_size=vocab_size, norm_type=norm_type, mlp_type=mlp_type)
g = torch.Generator().manual_seed(seed)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)
for step in range(steps):
x, y = p5.get_batch(train_data, cfg.block_size, 32, 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()
# Final eval
with torch.no_grad():
_, train_loss = model(*p5.get_batch(train_data, cfg.block_size, 32, g))
_, val_loss = model(*p5.get_batch(val_data, cfg.block_size, 32, g))
return model, float(train_loss.item()), float(val_loss.item()) # type: ignore[union-attr]
def count_params(model: ModernGPT) -> int:
seen = set()
total = 0
for p in model.parameters():
if id(p) not in seen:
seen.add(id(p))
total += p.numel()
return total
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 = 2000 if args.full else 200
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:]
configs = [
("ln", "gelu"),
("rms", "gelu"),
("ln", "swiglu"),
("rms", "swiglu"),
]
results = []
for norm_type, mlp_type in configs:
model, train_loss, val_loss = train_variant(
norm_type,
mlp_type,
vocab_size,
train_data,
val_data,
steps=args.steps,
seed=args.seed,
)
n_params = count_params(model)
results.append((norm_type, mlp_type, n_params, train_loss, val_loss))
print(
f" {norm_type:5s} + {mlp_type:8s} params={n_params:>6d} "
f"train={train_loss:.4f} val={val_loss:.4f}"
)
print(f"\nUniform baseline: {math.log(vocab_size):.4f}")
print(f"\n{'norm':6s} {'mlp':10s} {'params':>8s} {'train':>8s} {'val':>8s}")
print("-" * 50)
for norm_type, mlp_type, n_params, train_loss, val_loss in results:
print(
f"{norm_type:6s} {mlp_type:10s} {n_params:>8d} {train_loss:>8.4f} {val_loss:>8.4f}"
)
# === Outputs ===
log = args.output_dir / "run_log.txt"
lines = [
"# Project 7 run log",
f"steps: {args.steps}",
f"uniform baseline: {math.log(vocab_size):.4f}",
"",
f"{'norm':6s} {'mlp':10s} {'params':>8s} {'train':>8s} {'val':>8s}",
]
for norm_type, mlp_type, n_params, train_loss, val_loss in results:
lines.append(
f"{norm_type:6s} {mlp_type:10s} {n_params:>8d} {train_loss:>8.4f} {val_loss:>8.4f}"
)
lines.append("")
lines.append("Lessons:")
lines.append("- RMSNorm: same regularization effect as LayerNorm, ~half the FLOPs.")
lines.append("- SwiGLU: gating learns which features to pass through; modest extra")
lines.append(" params (we scaled hidden_mult=8/3 to roughly match GELU MLP size).")
log.write_text("\n".join(lines) + "\n", encoding="utf-8")
print(f"\nOutputs written to {args.output_dir.resolve()}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
The Steps
7 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 01Build Layernorm From Scratchstep_01_build-layernorm-from-scratch.py
"""
Project 7: Step 1 — Build LayerNorm from scratch
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.
"""
class LayerNorm(nn.Module):
def __init__(self, d_model, eps=1e-5):
super().__init__()
self.weight = nn.Parameter(torch.ones(d_model)) # gamma
self.bias = nn.Parameter(torch.zeros(d_model)) # beta
self.eps = eps
def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
x_hat = (x - mean) / torch.sqrt(var + self.eps)
return self.weight * x_hat + self.bias
Step 02Build Rmsnorm From Scratchstep_02_build-rmsnorm-from-scratch.py
"""
Project 7: Step 2 — Build RMSNorm from scratch
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.
"""
class RMSNorm(nn.Module):
def __init__(self, d_model, eps=1e-5):
super().__init__()
self.weight = nn.Parameter(torch.ones(d_model))
self.eps = eps
def forward(self, x):
rms = torch.sqrt((x ** 2).mean(dim=-1, keepdim=True) + self.eps)
x_hat = x / rms
return self.weight * x_hat
Step 03Wire Normalization Into The Transformer Blockstep_03_wire-normalization-into-the-transformer-block.py
"""
Project 7: Step 3 — Wire normalization into the transformer block
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.
"""
def make_norm(norm_type: str, d_model: int) -> nn.Module:
if norm_type == "layernorm":
return LayerNorm(d_model)
if norm_type == "rmsnorm":
return RMSNorm(d_model)
if norm_type == "none":
return nn.Identity()
raise ValueError(f"unknown norm_type: {norm_type}")
self.norm1 = make_norm(config.norm_type, config.d_model)
self.norm2 = make_norm(config.norm_type, config.d_model)
x = x + self.attn(self.norm1(x))
x = x + self.mlp(self.norm2(x))
Step 04Instrument Hidden State Statisticsstep_04_instrument-hidden-state-statistics.py
"""
Project 7: Step 4 — Instrument hidden-state statistics
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.
"""
def tensor_stats(x: torch.Tensor) -> dict:
flat = x.detach().float().reshape(-1)
mean = flat.mean()
var = flat.var(unbiased=False)
std = torch.sqrt(var + 1e-12)
centered = flat - mean
kurt = (centered ** 4).mean() / (std ** 4 + 1e-12)
return {
"mean": mean.item(),
"var": var.item(),
"min": flat.min().item(),
"max": flat.max().item(),
"kurtosis": kurt.item(),
}
if step % config.stats_every == 0:
stats["block_03_pre_attn"] = tensor_stats(x)
Step 06Build The Activation Functionsstep_06_build-the-activation-functions.py
"""
Project 7: Step 6 — Build the activation functions
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.
"""
def gelu(x):
return 0.5 * x * (1.0 + torch.tanh(
math.sqrt(2.0 / math.pi) * (x + 0.044715 * x**3)
))
def relu2(x):
return torch.relu(x) ** 2
class SwiGLU(nn.Module):
def __init__(self, d_model, hidden_dim):
super().__init__()
self.w1 = nn.Linear(d_model, hidden_dim, bias=False)
self.w2 = nn.Linear(d_model, hidden_dim, bias=False)
self.w3 = nn.Linear(hidden_dim, d_model, bias=False)
def forward(self, x):
gate = F.silu(self.w1(x))
value = self.w2(x)
return self.w3(gate * value)
Step 07Make The Mlp Activation Configurablestep_07_make-the-mlp-activation-configurable.py
"""
Project 7: Step 7 — Make the MLP activation configurable
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.
"""
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
hidden = 4 * config.d_model
if config.activation == "swiglu":
self.kind = "swiglu"
self.w1 = nn.Linear(config.d_model, hidden, bias=False)
self.w2 = nn.Linear(config.d_model, hidden, bias=False)
self.w3 = nn.Linear(hidden, config.d_model, bias=False)
else:
self.kind = config.activation
self.fc1 = nn.Linear(config.d_model, hidden)
self.fc2 = nn.Linear(hidden, config.d_model)
def forward(self, x):
if self.kind == "gelu":
return self.fc2(gelu(self.fc1(x)))
if self.kind == "relu2":
return self.fc2(torch.relu(self.fc1(x)) ** 2)
if self.kind == "swiglu":
return self.w3(F.silu(self.w1(x)) * self.w2(x))
raise ValueError(self.kind)
Step 08Replace Learned Positional Embeddings With Ropestep_08_replace-learned-positional-embeddings-with-rope.py
"""
Project 7: Step 8 — Replace learned positional embeddings with RoPE
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 = tok_emb[idx] + pos_emb[pos]
def rotate_half(x):
x1 = x[..., ::2]
x2 = x[..., 1::2]
return torch.stack((-x2, x1), dim=-1).flatten(-2)
def apply_rope(x, cos, sin):
return (x * cos) + (rotate_half(x) * sin)
def build_rope_cache(seq_len: int, head_dim: int, device, base: int = 10000) -> tuple:
half_dim = head_dim // 2
freq_seq = torch.arange(half_dim, device=device).float()
inv_freq = 1.0 / (base ** (freq_seq / half_dim))
positions = torch.arange(seq_len, device=device).float()
angles = torch.outer(positions, inv_freq) # (T, head_dim/2)
cos = torch.cos(angles).repeat_interleave(2, dim=-1) # (T, head_dim)
sin = torch.sin(angles).repeat_interleave(2, dim=-1)
return cos[None, None, :, :], sin[None, None, :, :]
cos, sin = build_rope_cache(T, head_dim, x.device)
q = apply_rope(q, cos, sin)
k = apply_rope(k, cos, sin)
Break It
break_it.pyThe four-variant comparison IS the experiment. If you want a "broken" version, run with norm_type="ln", mlp_type="gelu" (the older recipe) as a baseline and compare. The cleaner "what fails" exercise is to disable LayerNorm entirely (zero the gain parameter, freeze it) — left as an exercise.
$python projects/07_the-details-that-matter/break_it.py --tiny
"""
Project 7: BREAK IT experiment.
Deliberately sabotages one mechanism from build.py to show what happens
when it's removed. Compare outputs to the working version.
"""
self.norm1 = make_norm(config.norm_type, config.d_model)
self.norm2 = make_norm(config.norm_type, config.d_model)
def norm_for_layer(layer_idx, n_layers, d_model):
if layer_idx < n_layers // 2:
return LayerNorm(d_model)
return RMSNorm(d_model)
self.norm1 = norm_for_layer(layer_idx, config.n_layers, config.d_model)
self.norm2 = norm_for_layer(layer_idx, config.n_layers, config.d_model)
Outputs
run_log.txt
# Project 7 run log
steps: 200
uniform baseline: 3.8712
norm mlp params train val
ln gelu 105216 1.5537 2.6379
rms gelu 104896 1.5506 2.5884
ln swiglu 104320 1.5969 2.5978
rms swiglu 104000 1.6135 2.5991
Lessons:
- RMSNorm: same regularization effect as LayerNorm, ~half the FLOPs.
- SwiGLU: gating learns which features to pass through; modest extra
params (we scaled hidden_mult=8/3 to roughly match GELU MLP size).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 7 of Under the Hood.
Get the book →