What the optimizer holds

The chapter above this lesson calls opt.step() the line where parameters actually move. Which parameters is a question answered once, at construction: model.parameters() is a generator, and the optimizer drains it into a flat list of tensor references that nothing ever re-reads.

That sentence is easy to nod along to, so run the version that bites. Build the optimizer, then replace a layer on the model. The new weight is in the graph, so backward fills its .grad exactly as you would expect. It is not in the optimizer's list, so step() walks past it. Nothing raises, and one layer of the model quietly never learns.

Two operations exist for the case where the parameter set legitimately changes. add_param_group appends to the list, which is how a fine-tuning script unfreezes a backbone partway through a run. Rebuilding the optimizer is the other, and it throws away the state that the third section of this lesson is about.

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(8, 4), nn.ReLU(), nn.Linear(4, 1))
opt = torch.optim.SGD(model.parameters(), lr=0.1)

model[2] = nn.Linear(4, 1)        # a layer replaced after the optimizer was built
x = torch.randn(16, 8)
before = model[2].weight.detach().clone()

loss = nn.functional.mse_loss(model(x), x.sum(dim=1, keepdim=True))
opt.zero_grad()
loss.backward()
opt.step()

print(model[2].weight.grad is not None)               # True: it did get a gradient
print(torch.equal(before, model[2].weight.detach()))  # True: it never moved

param_groups is a list of plain dictionaries. Each one holds its own params list plus every hyperparameter the optimizer takes, filled in from the constructor defaults. An AdamW built with nothing but a learning rate still carries ten of them, fused and capturable included.

So the learning rate is not one number on the optimizer. It is one number per group, and anything that changes it changes it group by group. The third lesson in this arc lives on that fact.

Groups exist so one optimizer can apply different hyperparameters to different parameters, and the common use is weight decay: decay the weight matrices, leave the biases alone. Hand the constructor a list of dicts instead of a bare iterable and each dict becomes a group, inheriting whatever it does not override.

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 1))
opt = torch.optim.AdamW(model.parameters(), lr=1e-2)

print(len(opt.param_groups), len(opt.param_groups[0]["params"]))    # 1 4
print(sorted(k for k in opt.param_groups[0] if k != "params"))
# ['amsgrad', 'betas', 'capturable', 'differentiable', 'eps', 'foreach',
#  'fused', 'lr', 'maximize', 'weight_decay']

Construct an AdamW and len(opt.state) reads zero. The running estimates are allocated lazily, on the first step that sees a gradient for a given parameter, which is why an optimizer you built but never stepped serializes to an empty state.

The dictionary is keyed by the Parameter object, not by a name and not by an index. Ask for opt.state[model[0].weight] and you get the three tensors Adam keeps for that weight: exp_avg, exp_avg_sq, and step. That last one is a tensor rather than a Python int, which starts to matter the moment the state has to move to a device or be captured inside a graph.

Laziness has a practical edge worth carrying. A parameter that has never received a gradient has no state at all, so a checkpoint taken before the first step and one taken after are different objects, and both of them load without complaint.

print(len(opt.state))                       # 0: nothing has stepped yet

x = torch.randn(64, 8)
nn.functional.mse_loss(model(x), x.sum(dim=1, keepdim=True)).backward()
opt.step()

w = model[0].weight
print(len(opt.state), sorted(opt.state[w]))   # 4 ['exp_avg', 'exp_avg_sq', 'step']
print(opt.state[w]["step"])                   # tensor(1.)

opt.state_dict() cannot serialize tensor keys, so it replaces them with integers, counting through the groups in order. The state comes back keyed 0, 1, 2, 3, and each group's params entry becomes the list of indices belonging to it.

load_state_dict rebuilds the mapping by zipping the saved indices against the live parameters in the same order. Two checks run before that zip, and both are about counts: how many groups, and how many parameters in each group. The twelve lines below are the whole of it.

Read what is not there. No shape is compared. No name exists at this layer to compare. Reorder the parameters you hand the constructor and the checkpoint still loads, mapping Adam's memory for layer one onto layer three.

        if len(groups) != len(saved_groups):
            raise ValueError("loaded state dict has a different number of "
                             "parameter groups")
        param_lens = (len(g['params']) for g in groups)
        saved_lens = (len(g['params']) for g in saved_groups)
        if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)):
            raise ValueError("loaded state dict contains a parameter group "
                             "that doesn't match the size of optimizer's group")

        # Update the state
        id_map = dict(zip(chain.from_iterable(g['params'] for g in saved_groups),
                      chain.from_iterable(g['params'] for g in groups)))

Take the state dict from a model with a 32-wide hidden layer and load it into a model with a 4-wide one. Same number of groups, same four parameters, so both checks pass and the load returns cleanly. The optimizer now holds an exp_avg of shape (32, 8) for a parameter of shape (4, 8).

The failure arrives one step later, from inside AdamW, as an ordinary broadcasting error out of exp_avg.lerp_(grad, 1 - beta1). That is a good error to recognize on sight, because the line it names has nothing to do with the mistake, which happened at load time in a different function.

The other half of a checkpoint behaves the opposite way. model.state_dict() is keyed by parameter name and refuses a mismatch at load time, which the modules arc under chapter three demonstrates on a buffer key. Only one of the two halves tells you at the moment you were wrong.

import torch
from torch import nn

torch.manual_seed(0)
big = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 1))
opt = torch.optim.AdamW(big.parameters(), lr=1e-2)
x = torch.randn(64, 8)
nn.functional.mse_loss(big(x), x.sum(dim=1, keepdim=True)).backward()
opt.step()

sd = opt.state_dict()
print(list(sd), sorted(sd["state"]), sd["param_groups"][0]["params"])
# ['state', 'param_groups'] [0, 1, 2, 3] [0, 1, 2, 3]

small = nn.Sequential(nn.Linear(8, 4), nn.ReLU(), nn.Linear(4, 1))
opt2 = torch.optim.AdamW(small.parameters(), lr=1e-2)
opt2.load_state_dict(sd)                      # accepted, no complaint
print(opt2.state[small[0].weight]["exp_avg"].shape, small[0].weight.shape)
# torch.Size([32, 8]) torch.Size([4, 8])

nn.functional.mse_loss(small(x), x.sum(dim=1, keepdim=True)).backward()
opt2.step()
# RuntimeError: The size of tensor a (32) must match the size of tensor b (4)
# at non-singleton dimension 0

The autograd arc under chapter two covers what zero_grad does to a reference you were holding. The question here is what the optimizer does next, and it is a different one: every optimizer in torch.optim skips a parameter whose .grad is None outright, so a cleared parameter is not merely left un-updated, it is never visited.

A grad of zeros is a different instruction. The parameter is not skipped, the update runs, and an update with momentum or weight decay in it does not need a gradient to move anything. Momentum keeps discharging the buffer it built; weight decay keeps pulling toward zero.

The arithmetic is small enough to check by hand. SGD at lr 0.1, momentum 0.9, weight decay 0.1, one real step from 1.0 with grad 0.5 lands at 0.94 and leaves a buffer of 0.6. Step again on a zero grad and the effective gradient is 0 + 0.1 * 0.94, the buffer becomes 0.9 * 0.6 + 0.094, and the parameter lands at 0.8766. Under the default it would still read 0.94.

A parameter with no gradient this step and a parameter with a zero gradient this step are different requests, and only one of them holds still.

import torch


def one():
    p = torch.nn.Parameter(torch.ones(3))
    opt = torch.optim.SGD([p], lr=0.1, momentum=0.9, weight_decay=0.1)
    p.grad = torch.full((3,), 0.5)
    opt.step()                       # one real step, so momentum exists
    return p, opt


a, oa = one()
b, ob = one()
print(a[0].item(), b[0].item())      # 0.9399999976158142 0.9399999976158142

oa.zero_grad()                       # the default: set_to_none=True
ob.zero_grad(set_to_none=False)
print(a.grad, b.grad)                # None tensor([0., 0., 0.])

oa.step()
ob.step()
print(a[0].item(), b[0].item())      # 0.9399999976158142 0.8766000270843506

The modules arc establishes the model side of this: model.state_dict() hands back tensors sharing storage with the live parameters, so an in-memory snapshot moves when the model does. The optimizer side is a different mechanism, and it shows up one restore later than you would look for it.

Optimizer.load_state_dict casts each saved tensor to the target parameter's device and dtype, and when those already match, the cast returns the same tensor. The restored optimizer then writes its updates into the dictionary you loaded. Module.load_state_dict copies through param.copy_ and does not do this, so only one half of a checkpoint is consumed by restoring it.

So a checkpoint dict held in memory serves exactly one restore. A torch.load per restore is the fix, or a deepcopy if the dict has to stay in memory. It costs a file read, and it is the difference between three comparable resume experiments and three that quietly share state.

import torch
from torch import nn

torch.manual_seed(0)
m = nn.Linear(4, 2)
opt = torch.optim.AdamW(m.parameters(), lr=1e-2)
m(torch.randn(3, 4)).sum().backward()
opt.step()

ckpt = opt.state_dict()
kept = ckpt["state"][0]["exp_avg"].clone()

m2 = nn.Linear(4, 2)
opt2 = torch.optim.AdamW(m2.parameters(), lr=1e-2)
opt2.load_state_dict(ckpt)
print(opt2.state[m2.weight]["exp_avg"] is ckpt["state"][0]["exp_avg"])   # True
m2(torch.randn(3, 4)).sum().backward()
opt2.step()
print(torch.equal(kept, ckpt["state"][0]["exp_avg"]))   # False: the restore consumed it