Optimizers and Learning-Rate Schedules

Each group dict overrides defaults. state is keyed by the Parameter object and filled lazily on the first step(). For Linear(3, 4) → Linear(4, 2) after one AdamW step, optimizer.state_dict() printed (locally):

param_groups[0] = {'lr': 0.001, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 0.01,
  'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False,
  'differentiable': False, 'fused': None, 'decoupled_weight_decay': True, 'params': [0, 1, 2, 3]}
state = {0: {step: (), exp_avg: (4, 3), exp_avg_sq: (4, 3)}, 1: {... (4,)}, 2: {... (2, 4)}, 3: {... (2,)}}

Tensors become integer ids in iteration order, so load_state_dict matches state to parameters purely by position. Rebuild the optimizer from the same parameters in the same order and grouping, or Adam moments land on the wrong tensors, silently when shapes happen to match. Never pass a set. Passing named_parameters() stores param_names, but loading is still positional (optim docs). Note decoupled_weight_decay: True: in 2.8 AdamW is literally Adam(..., decoupled_weight_decay=True) (adamw.py at v2.8.0).

step() reads p.grad for every parameter, skips those whose grad is None, and updates p in place under torch.no_grad(), so the parameter stays a leaf (module 07). zero_grad(set_to_none=True) is the default since 2.0: it skips a memset and the next backward assigns instead of accumulating, as the Performance Tuning Guide (living page, accessed 2026-10-10) recommends. The semantic difference: a None grad is skipped, a zero grad is stepped, so with set_to_none=False a parameter unused this step is still decayed and still coasts on momentum.

Line-search methods need the loss several times per step, so step(closure) accepts a function that zeroes grads, recomputes the loss, calls backward() and returns it. LBFGS(lr=1, max_iter=20, history_size=100, line_search_fn=None) (docs) is for full-batch deterministic problems; locally, on a 2-D float64 quadratic with line_search_fn="strong_wolfe", one step called the closure 6 times and matched the solution of \(Ap = c\) to printed precision.

Hooks: register_step_pre_hook(hook(opt, args, kwargs)) and register_step_post_hook wrap every step() (log update norms); register_state_dict_pre_hook, register_load_state_dict_post_hook and their twins rewrite checkpoints, e.g. remapping state by param_names after a model change, as the optim docs show.

With lr \(\eta\), momentum \(\mu\), dampening \(\tau\) and weight decay \(\lambda\), torch.optim.SGD computes (sgd.py at v2.8.0):

With constant \(\eta\) the forms are identical (\(v = -\eta b\)). Under a schedule they are not: PyTorch reacts to an lr drop at once, while Sutskever's form coasts on velocity built at the old rate. Fence 1 shows both. Nesterov requires momentum > 0 and dampening == 0.

Adam (Kingma & Ba) keeps two moving averages and corrects their zero-initialization bias:

The two decays differ in where \(\lambda\theta\) enters. Adam(weight_decay=λ) is L2: \(g_t \leftarrow g_t + \lambda\theta_{t-1}\) before the moments, so the decay is divided by \(\sqrt{\hat v_t}\) like everything else. AdamW is decoupled (Loshchilov & Hutter): \(\theta \leftarrow \theta(1-\eta\lambda)\) before the Adam step, and \(m\), \(v\) never see it (AdamW docs).

Under L2 the effective decay of coordinate \(i\) is about \(\eta\lambda\theta_i/\sqrt{\hat v_i}\): parameters with small gradient variance are decayed far harder than requested, noisy ones barely. At zero gradient, \(g = \lambda\theta\), \(m/\sqrt{v} \to 1\), and Adam+L2 shrinks the weight by \(\eta\) per step whatever \(\lambda\) is: fence 2 measures 0.04995 over 5 steps at \(\eta = 10^{-2}\) for both \(\lambda = 0.1\) and \(0.01\), against 0.00499 for AdamW. Decoupled decay shrinks every weight by the same \(1-\eta\lambda\). PyTorch multiplies by the scheduled \(\eta\), so weight_decay=0.1 at lr 3e-3 shrinks weights by \(3\times10^{-4}\) per step, and the paper's \(\lambda\) equals \(\eta_{\text{peak}}\lambda_{\text{torch}}\).

amsgrad=True divides by the running maximum of \(v_t\). State memory: exp_avg and exp_avg_sq each match the parameter, so fp32 Adam holds 8 bytes per parameter of state (12 with amsgrad; SGD momentum 4). The state stays fp32 under autocast mixed precision, because the parameters do. Add 4-byte weights and 4-byte grads and you reach the 16 bytes/parameter of module 22.

Adafactor (Shazeer & Stern; in torch.optim since 2.5, docs) factors the second moment of a matrix: locally a 4096×1024 weight got row_var (4096, 1) and col_var (1, 1024), 5,121 state numbers against AdamW's 8,388,609, \(O(m+n)\) instead of \(O(mn)\), and no first moment. Dense optimizers reject sparse gradients. Muon, Shampoo, SOAP and Lion are not in torch.optim 2.8; see the optimization notes.

Per the optim docs, only Adam, AdamW, SGD and Adagrad have fused kernels (Adam/AdamW stable on CUDA, beta on CPU and MPS; Adagrad CPU-only). Locally fused=True ran on CPU for those four and raised TypeError for RMSprop, NAdam and Adafactor. Implementations agree within float tolerance, not guaranteed bitwise: summation order and fused multiply-adds differ.

capturable=True keeps step and lr on device for CUDA-graph capture (not executed here); differentiable=True keeps the update in the autograd graph for meta-learning; maximize=True ascends. lr may be a tensor: in a local probe with a counting torch.compile backend, changing a float lr between steps compiled the step twice, changing a tensor lr's value compiled it once.

Decay says matrices that multiply activations should be small. It makes little sense for biases (a shift) or LayerNorm gains (a scale initialized at 1, see module 13, that decay drags toward 0). nanoGPT's configure_optimizers (model.py at 3adf61e) groups by rank, which the lab copies:

def build_param_groups(model, weight_decay):
    decay, no_decay, seen = [], [], set()
    for _, p in model.named_parameters():      # tied tensors appear once
        if not p.requires_grad or id(p) in seen:
            continue
        seen.add(id(p))
        (decay if p.ndim >= 2 else no_decay).append(p)
    return [{"params": decay, "weight_decay": weight_decay},
            {"params": no_decay, "weight_decay": 0.0}]

ndim >= 2 catches every Linear and Embedding weight. Decaying embeddings is a choice: nanoGPT does (they are 2-D and tied to the head); many recipes exclude them. named_parameters() deduplicates tied weights; the id set guards hand-built lists. A tensor in two groups raises ValueError: some parameters appear in more than one parameter group; a duplicate inside one group only warns in 2.8, and its update is then applied twice (measured).

Groups also carry learning rates: a smaller lr for a pretrained backbone, or layer-wise decay \(\eta_\ell = \eta\,\gamma^{L-\ell}\) (module 19, module 36). Schedulers scale each group's own initial_lr, so ratios survive. opt.add_param_group({"params": backbone.parameters(), "lr": 1e-5}) unfreezes later, with empty state.

A scheduler keeps last_epoch (a step counter, despite the name) and rewrites group["lr"] on each scheduler.step(). The constructor itself performs step 0, so with LambdaLR(opt, f) the first optimizer step already uses \(\eta_0 f(0)\), not the bare lr you passed to the optimizer (measured).

OneCycleLR (Smith, docs) also cycles momentum: on AdamW, betas[0] fell from 0.95 to 0.85 as lr rose (measured), and stepping past total_steps raises ValueError.

Order: optimizer.step() then scheduler.step(). Reversed, PyTorch warns ("Detected call of lr_scheduler.step() before optimizer.step()") and the first schedule value is skipped. LambdaLR cannot tell iterations from epochs: give the lambda total steps and step it every iteration.

The capstone (module 33) warms up linearly to 3e-3 over \(W=40\) steps, then cosine-decays to a \(\rho = 10\,\%\) floor at \(T=800\):

As LambdaLR this is one function (docs). The built-in composition is SequentialLR([LinearLR(start_factor=1/40, total_iters=39), CosineAnnealingLR(T_max=760, eta_min=3e-4)], milestones=[40]). total_iters=39, not 40: LinearLR reaches end_factor at t = total_iters, and the peak belongs at step 39 (fence 4). Prefer the lambda: nothing to misalign.

Schedulers are stateful: checkpoint scheduler.state_dict() (last_epoch, base_lrs, _last_lr, ...). Rebuilding it on resume restarts warmup (fence 3, module 16).

AveragedModel(model, multi_avg_fn=get_ema_multi_avg_fn(0.999)) from torch.optim.swa_utils holds a deep copy of the model; call ema.update_parameters(model) after each optimizer.step(). The update is \(W^{\text{EMA}} \leftarrow \alpha W^{\text{EMA}} + (1-\alpha)W\) via torch._foreach_lerp_, and the first call copies the weights outright. The horizon is about \(1/(1-\alpha)\) steps (1,000 for 0.999), so the lab's 200-step run uses 0.95. Without multi_avg_fn it computes the equal-weight SWA average (Izmailov et al.), usually with SWALR(optimizer, swa_lr=...), which anneals to a constant lr.

Buffers are not averaged unless use_buffers=True, so BatchNorm models need update_bn(loader, swa_model) to recompute running statistics. Use EMA weights for evaluation and generation; training continues on the raw weights. The copy costs another full set of parameters (4 bytes/parameter in fp32).

Put hyperparameters in defaults, iterate self.param_groups, decorate step with @torch.no_grad(), create self.state[p] lazily, support a closure. Keep everything in self.state and the group dicts and state_dict() round-trips for free. The core of the lab's Lion-like SignMomentum:

@torch.no_grad()
def step(self, closure=None):
    loss = None
    if closure is not None:
        with torch.enable_grad():
            loss = closure()
    for group in self.param_groups:
        lr, (b1, b2), wd = group["lr"], group["betas"], group["weight_decay"]
        for p in group["params"]:
            if p.grad is None:
                continue
            state = self.state[p]
            if "exp_avg" not in state:                            # lazy state
                state["exp_avg"] = torch.zeros_like(p)
            m = state["exp_avg"]
            p.mul_(1 - lr * wd)                                   # decoupled decay
            p.add_(torch.sign(m.mul(b1).add(p.grad, alpha=1 - b1)), alpha=-lr)
            m.mul_(b2).add_(p.grad, alpha=1 - b2)
    return loss

Every coordinate moves by exactly \(\pm\eta\), so its lr must be several times smaller than AdamW's; it stores one state tensor.

1. SGD momentum and Nesterov by hand. Proves that PyTorch's buffer is \(b \leftarrow \mu b + g\) with \(b_1 = g_1\), that Nesterov steps along \(g + \mu b\), and that Sutskever's velocity form agrees under a constant lr but diverges after an lr drop. It does not test weight decay (the lab does, at 1e-4) or dampening.

import torch

A = torch.tensor([[3.0, 0.5], [0.5, 1.0]], dtype=torch.float64)
c = torch.tensor([1.0, -2.0], dtype=torch.float64)
grad = lambda p: A @ p - c                     # f(p) = 0.5 p^T A p - c^T p
lr, mu = 0.1, 0.9

for nesterov in (False, True):
    p = torch.nn.Parameter(torch.ones(2, dtype=torch.float64))
    opt = torch.optim.SGD([p], lr=lr, momentum=mu, nesterov=nesterov)
    q, b = torch.ones(2, dtype=torch.float64), None
    for _ in range(5):
        opt.zero_grad()
        (0.5 * p @ A @ p - c @ p).backward()
        opt.step()
        g = grad(q)
        b = g.clone() if b is None else mu * b + g    # step 1: b = g
        q = q - lr * (g + mu * b if nesterov else b)
        torch.testing.assert_close(p.detach(), q, rtol=0, atol=1e-14)

def pytorch_form(lrs):                         # b = mu b + g ; p -= lr_t b
    p, b = torch.ones(2, dtype=torch.float64), None
    for lr_t in lrs:
        b = grad(p) if b is None else mu * b + grad(p)
        p = p - lr_t * b
    return p

def sutskever_form(lrs):                       # v = mu v - lr_t g ; p += v
    p, v = torch.ones(2, dtype=torch.float64), 0.0
    for lr_t in lrs:
        v = mu * v - lr_t * grad(p)
        p = p + v
    return p

const, drop = [0.1] * 6, [0.1] * 3 + [0.01] * 3
torch.testing.assert_close(pytorch_form(const), sutskever_form(const), rtol=0, atol=1e-14)
gap = (pytorch_form(drop) - sutskever_form(drop)).abs().max().item()
assert gap > 1e-2
print(f"SGD momentum/Nesterov match by hand; after an lr drop the two forms differ by {gap:.3f}")

2. Adam vs AdamW differ only through decay. Proves equality at weight_decay=0, that AdamW is Adam(decoupled_weight_decay=True), that AdamW matches the hand-written decoupled update to 1e-12 in float64, and the zero-gradient pathology of L2 inside Adam. Gradients are a fixed sequence, so it checks the update rule, not training dynamics.

import torch

torch.manual_seed(0)
w0 = torch.randn(4, 3, dtype=torch.float64)
grads = [torch.randn(4, 3, dtype=torch.float64) for _ in range(5)]
lr, wd, (b1, b2), eps = 1e-2, 0.1, (0.9, 0.999), 1e-8
Adam, AdamW = torch.optim.Adam, torch.optim.AdamW

def run(cls, w, gs, **kw):
    p = torch.nn.Parameter(w.clone())
    opt = cls([p], lr=lr, **kw)
    for g in gs:
        p.grad = g.clone()
        opt.step()
    return p.detach()

adamw = run(AdamW, w0, grads, weight_decay=wd)
assert torch.equal(run(Adam, w0, grads), run(AdamW, w0, grads, weight_decay=0.0))
assert torch.equal(run(Adam, w0, grads, weight_decay=wd, decoupled_weight_decay=True), adamw)
assert not torch.allclose(run(Adam, w0, grads, weight_decay=wd), adamw)

p, m, v = w0.clone(), torch.zeros_like(w0), torch.zeros_like(w0)
for t, g in enumerate(grads, start=1):        # AdamW by hand
    p = p * (1 - lr * wd)                      # decoupled decay first
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g * g
    p = p - lr * (m / (1 - b1 ** t)) / ((v / (1 - b2 ** t)).sqrt() + eps)
torch.testing.assert_close(adamw, p, rtol=0, atol=1e-12)

one, zero = torch.ones(1, dtype=torch.float64), [torch.zeros(1, dtype=torch.float64)] * 5
shrink = lambda cls, d: 1 - run(cls, one, zero, weight_decay=d).item()
assert abs(shrink(AdamW, wd) - (1 - (1 - lr * wd) ** 5)) < 1e-12
assert abs(shrink(Adam, wd) - 5 * lr) < 1e-3 and abs(shrink(Adam, wd / 10) - 5 * lr) < 1e-3
print(f"zero-grad shrink over 5 steps: AdamW {shrink(AdamW, wd):.5f}, "
      f"Adam+L2 {shrink(Adam, wd):.5f} (wd/10: {shrink(Adam, wd / 10):.5f})")

3. state_dict round trip. Proves that optimizer plus scheduler state restores a bit-identical continuation through torch.save and torch.load(weights_only=True), that state is keyed by integer ids, and that a freshly built scheduler resets the lr to the warmup start. RNG and data-loader state are not covered (module 16).

import io
import torch
from torch import nn
from torch.optim.lr_scheduler import LambdaLR

torch.manual_seed(0)
x, y = torch.randn(16, 3), torch.randn(16, 1)
lam = lambda step: min(1.0, (step + 1) / 4)        # 4-step warmup

def build():
    model = nn.Sequential(nn.Linear(3, 4), nn.Tanh(), nn.Linear(4, 1))
    opt = torch.optim.AdamW(model.parameters(), lr=1e-2, betas=(0.9, 0.95))
    return model, opt, LambdaLR(opt, lam)

def train_step(model, opt, sched):
    opt.zero_grad(set_to_none=True)
    nn.functional.mse_loss(model(x), y).backward()
    opt.step()
    sched.step()

model, opt, sched = build()
for _ in range(3):
    train_step(model, opt, sched)
buf = io.BytesIO()
torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "sched": sched.state_dict()}, buf)
buf.seek(0)
ck = torch.load(buf, weights_only=True)
assert ck["opt"]["param_groups"][0]["params"] == [0, 1, 2, 3]        # ids, not tensors
assert sorted(ck["opt"]["state"][0]) == ["exp_avg", "exp_avg_sq", "step"]

model2, opt2, sched2 = build()
model2.load_state_dict(ck["model"])
opt2.load_state_dict(ck["opt"])
sched2.load_state_dict(ck["sched"])
_, opt3, _ = build()                                # the resume bug: fresh scheduler
opt3.load_state_dict(ck["opt"])
LambdaLR(opt3, lam)
assert abs(opt2.param_groups[0]["lr"] - 1e-2) < 1e-15 and abs(opt3.param_groups[0]["lr"] - 2.5e-3) < 1e-15

for _ in range(2):
    train_step(model, opt, sched)
    train_step(model2, opt2, sched2)
assert all(torch.equal(a, b) for a, b in zip(model.parameters(), model2.parameters()))
print("optimizer+scheduler round trip is bit-exact; a rebuilt scheduler restarts warmup at 2.5e-3")

4. Warmup-cosine values. Proves the capstone LambdaLR hits the closed form at steps 0, 39, 40, 420 and 800, that SequentialLR(LinearLR, CosineAnnealingLR) reproduces it within 1e-9 with total_iters=W-1, and that the natural-looking total_iters=W never reaches the peak.

import math
import warnings
import torch
from torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR, LinearLR, SequentialLR

peak, W, T, floor = 3e-3, 40, 800, 0.1

def mult(t):                                       # capstone schedule (module 33)
    if t < W:
        return (t + 1) / W
    return floor + (1 - floor) * 0.5 * (1 + math.cos(math.pi * min(1.0, (t - W) / (T - W))))

def trace(make_sched):
    opt = torch.optim.SGD([torch.nn.Parameter(torch.zeros(1))], lr=peak)
    sched, lrs = make_sched(opt), []
    for _ in range(T + 1):
        lrs.append(opt.param_groups[0]["lr"])      # lr used by this step
        opt.step()
        sched.step()
    return lrs

def sequential(linear_iters):
    return lambda opt: SequentialLR(opt, [
        LinearLR(opt, start_factor=1 / W, end_factor=1.0, total_iters=linear_iters),
        CosineAnnealingLR(opt, T_max=T - W, eta_min=floor * peak)], milestones=[W])

lam = trace(lambda opt: LambdaLR(opt, mult))
mid = (W + T) // 2
assert abs(lam[0] - peak / W) < 1e-15 and abs(lam[W - 1] - peak) < 1e-15 and abs(lam[W] - peak) < 1e-15
assert abs(lam[mid] - (floor + (1 - floor) / 2) * peak) < 1e-15
assert abs(lam[T] - floor * peak) < 1e-15

with warnings.catch_warnings():                    # SequentialLR passes an epoch at milestones
    warnings.simplefilter("ignore", UserWarning)
    seq, seq_naive = trace(sequential(W - 1)), trace(sequential(W))
assert max(abs(a - b) for a, b in zip(seq, lam)) < 1e-9
assert abs(seq_naive[W - 1] - peak) > 1e-6         # total_iters=W misses the peak
print(f"warmup-cosine lr[0]={lam[0]:.2e} lr[39]={lam[39]:.2e} lr[{mid}]={lam[mid]:.3e} lr[800]={lam[800]:.2e}")

code/15_optimizers_from_scratch.py implements ScratchSGD, ScratchAdamW and SignMomentum as Optimizer subclasses plus warmup_cosine(step), and trains a seeded 8 → 32 → LayerNorm → tanh → 1 MLP (385 parameters) on synthetic regression (256 train, 256 held-out rows): (a) scratch vs torch.optim SGD (momentum 0.9, with and without Nesterov, wd 1e-4) and AdamW (betas (0.9, 0.95), wd 0.1 on 2-D weights), float64, in lockstep under one LambdaLR, tracking the worst parameter gap at every step; (b) AdamW for-loop vs foreach=True vs fused=True in float32; (c) decay groups vs decay-all; (d) an EMA (decay 0.95) of the grouped run on held-out rows.

python content/courses/pytorch/code/15_optimizers_from_scratch.py
python site/test_pytorch_15_optimizers.py

Tested with Python 3.11 and PyTorch 2.8.0 on CPU.

One local PyTorch 2.8 CPU run (Apple M-series laptop, 200 steps, seed 0) produced:

The bitwise AdamW match comes from mirroring PyTorch's op order (lerp_ for \(m\), bias correction folded into the addcdiv_ step size). A variant writing \(m\) as mul_ then add_ drifted to 1.6e-15 after 200 steps: same math, different rounding. foreach on CPU falls back to per-tensor kernels, so it was bitwise equal here; on CUDA expect float32-epsilon differences, hence the test's < 1e-5. Decay-all shrank biases and gains and cost a little training loss. The EMA's edge is one seed, not evidence. The test asserts contracts: gaps below 1e-12, foreach and fused gaps below 1e-5, loss at least halved, scratch state_dict round trips, biases and gains untouched by a zero-gradient AdamW step while matrices shrink by exactly \(1-\eta\lambda\), tied weights listed once, and the lr trace equal to the closed form.

Scope, honestly: this shows equivalence of update rules on a toy problem. It says nothing about which optimizer generalizes better, and the fused CUDA kernels and capturable paths were not executed.