AdamW is the default optimizer for training transformers. This notebook builds it step by step, showing what each component does and why.
We’ll implement: 1. SGD — the simplest optimizer (just follow the gradient) 2. Momentum — remember past gradients for smoother updates 3. Adam — adaptive learning rates per parameter 4. AdamW — fix weight decay to be decoupled from gradients 5. Comparison — train with each and see the difference
src/ai_playground/training/optimizers.pythe backprop chapters of Neural Networks: Zero to Hero (Karpathy) build the gradient intuition behind these update rules — more in Videos.
Implement AdamW yourself — moments, bias correction, decoupled weight decay — until the update rule has no mystery left.
🏗️ Chapter milestone: M8 — prove your from-scratch optimizer on your own model — the exercise at the end of this chapter; see The Project.
One session, ending in 🏗️ M8 — prove your from-scratch optimizer on your model. Cards, prerequisites, and done-when tests: the Session Guide.
# Setup — works from a repo checkout (via ../src) or standalone on Colab
import sys
sys.path.insert(0, '../src')
try:
import ai_playground # noqa: F401
except ImportError:
%pip install -q git+https://github.com/ggreg/ai_playground.git
%config InlineBackend.figure_format = 'retina'
from ai_playground.plotting import apply_plot_style
apply_plot_style()import sys
sys.path.insert(0, '../src')
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import math
from ai_playground.models import Transformer, TransformerConfig
from ai_playground.training.optimizers import AdamWFromScratch
torch.manual_seed(42)
device = 'cpu'Stochastic Gradient Descent: move parameters in the opposite direction of the gradient.
θ = θ - lr × gradientSimple, but has two problems: - Same LR for all parameters — some weights need big steps, others small - No memory — doesn’t learn from past gradients, easily zig-zags
# Implement SGD from scratch
class SGDFromScratch:
def __init__(self, params, lr=0.01):
self.params = list(params)
self.lr = lr
def step(self):
with torch.no_grad():
for p in self.params:
if p.grad is not None:
# The entire algorithm: move opposite to gradient
p -= self.lr * p.grad
def zero_grad(self):
for p in self.params:
if p.grad is not None:
p.grad.zero_()
print('SGD: θ = θ - lr × ∇loss')
print('One line of math. Simple but slow to converge.')SGD: θ = θ - lr × ∇loss
One line of math. Simple but slow to converge.Instead of following only the current gradient, accumulate a running average of past gradients. This smooths out noise and accelerates through consistent directions (like a ball rolling downhill).
m = β × m + (1 - β) × gradient # exponential moving average
θ = θ - lr × mTypical β = 0.9 (remember ~10 past gradients).
class SGDMomentumFromScratch:
def __init__(self, params, lr=0.01, beta=0.9):
self.params = list(params)
self.lr = lr
self.beta = beta
self.m = [torch.zeros_like(p) for p in self.params] # momentum buffer
def step(self):
with torch.no_grad():
for p, m in zip(self.params, self.m):
if p.grad is not None:
# Update momentum: weighted average of past and current gradient
m.mul_(self.beta).add_(p.grad, alpha=1 - self.beta)
# Step in the smoothed direction
p -= self.lr * m
def zero_grad(self):
for p in self.params:
if p.grad is not None:
p.grad.zero_()
print('SGD + Momentum: m = β×m + (1-β)×∇loss, θ = θ - lr×m')
print('Smooths noisy gradients and accelerates in consistent directions.')SGD + Momentum: m = β×m + (1-β)×∇loss, θ = θ - lr×m
Smooths noisy gradients and accelerates in consistent directions.Different parameters need different learning rates. Adam tracks two things:
Parameters with large gradients get smaller effective LR. Parameters with small gradients get larger effective LR. Self-tuning!
m = β1 × m + (1 - β1) × g # momentum (direction)
v = β2 × v + (1 - β2) × g² # RMS of gradients (scale)
θ = θ - lr × m / (√v + ε) # adaptive step sizeAdam paper (Kingma & Ba, 2015)
class AdamFromScratch:
"""Adam without weight decay — for comparison with AdamW."""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8):
self.params = list(params)
self.lr = lr
self.beta1, self.beta2 = betas
self.eps = eps
self.t = 0
self.m = [torch.zeros_like(p) for p in self.params] # first moment
self.v = [torch.zeros_like(p) for p in self.params] # second moment
def step(self):
self.t += 1
with torch.no_grad():
for p, m, v in zip(self.params, self.m, self.v):
if p.grad is None:
continue
g = p.grad
# Update moments
m.mul_(self.beta1).add_(g, alpha=1 - self.beta1)
v.mul_(self.beta2).addcmul_(g, g, value=1 - self.beta2)
# Bias correction (moments start at 0, so they're
# underestimated in early steps)
m_hat = m / (1 - self.beta1 ** self.t)
v_hat = v / (1 - self.beta2 ** self.t)
# Adaptive update: large v → small step, small v → large step
p.addcdiv_(m_hat, v_hat.sqrt().add_(self.eps), value=-self.lr)
def zero_grad(self):
for p in self.params:
if p.grad is not None:
p.grad.zero_()
print('Adam: m = β1×m + (1-β1)×g, v = β2×v + (1-β2)×g²')
print(' θ = θ - lr × m̂ / (√v̂ + ε)')
print('Each parameter gets its own effective learning rate.')Adam: m = β1×m + (1-β1)×g, v = β2×v + (1-β2)×g²
θ = θ - lr × m̂ / (√v̂ + ε)
Each parameter gets its own effective learning rate.Weight decay prevents overfitting by shrinking parameters toward zero. Original Adam applied decay through the gradient (L2 regularization), but this interacts badly with adaptive learning rates.
AdamW fixes this by applying weight decay directly to parameters, independent of the gradient:
θ = θ × (1 - lr × λ) # shrink weights (decoupled!)
θ = θ - lr × m̂ / (√v̂ + ε) # then do the Adam updateThis is the standard optimizer for all modern LLM training.
See src/ai_playground/training/optimizers.py for the full implementation.
# Show our from-scratch AdamW matches PyTorch's
torch.manual_seed(42)
cfg = TransformerConfig(
vocab_size=8, dim=32, n_layers=2, n_heads=4,
n_kv_heads=2, max_seq_len=16, ffn_dim_multiplier=2.0,
)
# Model 1: PyTorch AdamW
model_pt = Transformer(cfg)
opt_pt = torch.optim.AdamW(model_pt.parameters(), lr=1e-3, weight_decay=0.01)
# Model 2: Our from-scratch AdamW
torch.manual_seed(42)
model_ours = Transformer(cfg)
opt_ours = AdamWFromScratch(model_ours.parameters(), lr=1e-3, weight_decay=0.01)
# One training step on same data
VOCAB = ['<pad>', 'the', 'cat', 'dog', 'sat', 'ran', 'on', 'mat']
tok2id = {t: i for i, t in enumerate(VOCAB)}
train_x = torch.tensor([[tok2id[w] for w in 'the cat sat on the'.split()]])
train_y = torch.tensor([[tok2id[w] for w in 'cat sat on the mat'.split()]])
for name, model, opt in [('PyTorch', model_pt, opt_pt), ('Ours', model_ours, opt_ours)]:
logits = model(train_x)
loss = F.cross_entropy(logits.view(-1, 8), train_y.view(-1))
opt.zero_grad()
loss.backward()
opt.step()
# Compare weights after one step
max_diff = max(
(p1 - p2).abs().max().item()
for p1, p2 in zip(model_pt.parameters(), model_ours.parameters())
)
print(f'Max weight difference after 1 step: {max_diff:.10f}')
print(f'Match: {max_diff < 1e-5} ✅' if max_diff < 1e-5 else f'Mismatch! ❌')Max weight difference after 1 step: 0.0000000075
Match: True ✅Let’s train the same model with SGD, SGD+Momentum, Adam, and AdamW and compare convergence.
sentences = ['the cat sat on the mat', 'the dog ran on the mat',
'the cat ran on the mat', 'the dog sat on the mat']
all_tok = [[tok2id[w] for w in s.split()] for s in sentences]
data_x = torch.tensor([t[:-1] for t in all_tok])
data_y = torch.tensor([t[1:] for t in all_tok])
def train_with_optimizer(make_opt_fn, name, n_steps=200, seed=42):
torch.manual_seed(seed)
model = Transformer(cfg)
opt = make_opt_fn(model.parameters())
losses = []
model.train()
for step in range(n_steps):
logits = model(data_x)
loss = F.cross_entropy(logits.view(-1, 8), data_y.view(-1))
if hasattr(opt, 'zero_grad'):
opt.zero_grad()
else:
for p in model.parameters():
if p.grad is not None:
p.grad.zero_()
loss.backward()
opt.step()
losses.append(loss.item())
print(f'{name:<25s} final loss: {losses[-1]:.4f}')
return losses
optimizers = [
('SGD (lr=0.1)', lambda p: SGDFromScratch(p, lr=0.1), '#94A3B8'),
('SGD+Momentum (lr=0.1)', lambda p: SGDMomentumFromScratch(p, lr=0.1), '#F97316'),
('Adam (lr=1e-3)', lambda p: AdamFromScratch(p, lr=1e-3), '#FF6B6B'),
('AdamW (lr=1e-3)', lambda p: AdamWFromScratch(p, lr=1e-3, weight_decay=0.01), '#4ECDC4'),
]
fig, ax = plt.subplots(figsize=(14, 6))
for name, make_fn, color in optimizers:
losses = train_with_optimizer(make_fn, name)
ax.plot(losses, label=name, color=color, linewidth=2)
ax.axhline(y=np.log(8), color='#CBD5E1', linestyle=':', label='Random guess')
ax.set_xlabel('Training Step', fontsize=12)
ax.set_ylabel('Loss', fontsize=12)
ax.set_title('Optimizer Comparison on Tiny Transformer', fontsize=14, fontweight='bold')
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()SGD (lr=0.1) final loss: 0.2833
SGD+Momentum (lr=0.1) final loss: 0.2815
Adam (lr=1e-3) final loss: 0.3153
AdamW (lr=1e-3) final loss: 0.3156
Let’s peek inside Adam’s state during training to see the first moment (m) and second moment (v) for one parameter.
# Train and record Adam's internal state
torch.manual_seed(42)
model = Transformer(cfg)
opt = AdamWFromScratch(model.parameters(), lr=1e-3, weight_decay=0.01)
# Track the first parameter of the embedding
target_param = model.tok_embeddings.weight
gradient_history = []
m_history = []
v_history = []
lr_effective = []
model.train()
for step in range(200):
logits = model(data_x)
loss = F.cross_entropy(logits.view(-1, 8), data_y.view(-1))
opt.zero_grad()
loss.backward()
# Record state before step
g = target_param.grad[1, 0].item() # "the" embedding, dim 0
gradient_history.append(g)
opt.step()
# Read Adam's state for this parameter
state = opt.state[target_param]
m_val = state['m'][1, 0].item()
v_val = state['v'][1, 0].item()
t = state['step']
# Bias-corrected values
m_hat = m_val / (1 - 0.9 ** t)
v_hat = v_val / (1 - 0.999 ** t)
eff_lr = 1e-3 * abs(m_hat) / (math.sqrt(v_hat) + 1e-8) if v_hat > 0 else 0
m_history.append(m_hat)
v_history.append(math.sqrt(v_hat))
lr_effective.append(eff_lr)
fig, axes = plt.subplots(4, 1, figsize=(14, 12), sharex=True)
axes[0].plot(gradient_history, color='#94A3B8', linewidth=1, alpha=0.7)
axes[0].set_ylabel('Raw gradient', fontsize=11)
axes[0].set_title('AdamW Internals for "the" embedding[0]', fontsize=14, fontweight='bold')
axes[0].grid(True, alpha=0.3)
axes[1].plot(m_history, color='#FF6B6B', linewidth=2)
axes[1].set_ylabel('m̂ (first moment)', fontsize=11)
axes[1].set_title('Smoothed gradient direction (momentum)', fontsize=12)
axes[1].grid(True, alpha=0.3)
axes[2].plot(v_history, color='#F97316', linewidth=2)
axes[2].set_ylabel('√v̂ (second moment)', fontsize=11)
axes[2].set_title('Gradient magnitude (for adaptive LR)', fontsize=12)
axes[2].grid(True, alpha=0.3)
axes[3].plot(lr_effective, color='#4ECDC4', linewidth=2)
axes[3].set_ylabel('Effective step size', fontsize=11)
axes[3].set_xlabel('Training Step', fontsize=12)
axes[3].set_title('lr × |m̂| / (√v̂ + ε) — how far this parameter actually moves', fontsize=12)
axes[3].grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
print('- Raw gradient is noisy → m̂ (momentum) smooths it')
print('- √v̂ tracks gradient scale → dividing by it normalizes the step size')
print('- Effective step size adapts per-parameter: big gradients → smaller steps')
- Raw gradient is noisy → m̂ (momentum) smooths it
- √v̂ tracks gradient scale → dividing by it normalizes the step size
- Effective step size adapts per-parameter: big gradients → smaller stepsRun 20 training steps of your model twice from identical init and data order: once with AdamWFromScratch, once with torch.optim.AdamW (same hyperparameters, fp32). Compare every parameter tensor:
worst = max((a - b).abs().max().item()
for a, b in zip(m1.parameters(), m2.parameters()))
assert worst < 1e-5, worstAccept when: the assert passes on your model — you now understand every update your release run made. Part of The Project: Serve Your Own LLM.
Session S2.8 — record the worst deviation in metrics.json['m8_optimizer_max_diff']; verify: uv run pytest tests/milestones/test_m08_adamw.py.
optimizers.py:70-85.00_training_from_scratch.ipynbThis chapter built AdamW with the repo’s code in front of you; M8 proved it on your model. The phase’s from-scratch project takes the reference away: your own AdamW (matched to torch.optim.AdamW to 1e-5), your own warmup+cosine schedule, and gradient accumulation proven exactly equivalent to the full batch — including the mean-loss rescaling trap this phase kept warning you about.
Brief and rules: projects/p2_trainer/README.md
uv run pytest projects/p2_trainer/ -v