The resume that matches

Constructing a scheduler is not a passive act. Before any step runs, it walks the optimizer's param groups and writes an initial_lr key into each one, then copies those values into its own base_lrs. Every learning rate it ever produces is computed from base_lrs and a counter, not from whatever lr currently reads.

That is why a scheduler needs a state dict of its own. Its whole memory is last_epoch and base_lrs plus its own hyperparameters, eight keys for a cosine schedule, and none of that lives in the optimizer.

It also explains a class of confusion around resuming. Set lr by hand on a param group, then let a scheduler step, and your value is gone: the scheduler recomputed the rate from base_lrs and overwrote it.

import torch
from torch.optim.lr_scheduler import CosineAnnealingLR

p = torch.nn.Parameter(torch.zeros(1))
opt = torch.optim.SGD([p], lr=0.1)
print("initial_lr" in opt.param_groups[0])                            # False

sched = CosineAnnealingLR(opt, T_max=10)
print("initial_lr" in opt.param_groups[0], opt.param_groups[0]["initial_lr"])
# True 0.1
print(sched.state_dict())
# {'T_max': 10, 'eta_min': 0, 'base_lrs': [0.1], 'last_epoch': 0,
#  'verbose': False, '_step_count': 1, '_get_lr_called_within_step': False,
#  '_last_lr': [0.1]}

Warmup then decay is two schedules on one optimizer, and SequentialLR is the object that hands over between them. Give it the schedulers, give it the step index where the second takes over, and step it once per optimizer step.

The ten learning rates a three-step linear warmup and a seven-step cosine actually produce are worth reading as a sequence rather than a formula. The warmup climbs 0.01, 0.04, 0.07 and hits the base rate of 0.1 on step four. The cosine then descends 0.095, 0.081, 0.061, 0.039, 0.019, 0.005.

SequentialLR keeps its own last_epoch and a nested list of its children's state dicts, so restoring it restores the hand-over point too. Composition here means one object that owns the others, not two objects racing to write the same field.

import torch
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR

p = torch.nn.Parameter(torch.zeros(1))
opt = torch.optim.SGD([p], lr=0.1)
sched = SequentialLR(
    opt,
    schedulers=[LinearLR(opt, start_factor=0.1, total_iters=3), CosineAnnealingLR(opt, T_max=7)],
    milestones=[3],
)

seen = []
for _ in range(10):
    opt.step()
    seen.append(round(opt.param_groups[0]["lr"], 6))
    sched.step()
print(seen)
# [0.01, 0.04, 0.07, 0.1, 0.095048, 0.081174, 0.061126, 0.038874, 0.018826,
#  0.004952]
print(sorted(sched.state_dict()), sched.state_dict()["last_epoch"])
# ['_last_lr', '_milestones', '_schedulers', 'last_epoch'] 10

The hand-over is six lines, and the interesting one is the branch. On the step where a milestone is reached, the incoming scheduler is stepped with an explicit epoch of 0, which restarts it at its own beginning rather than continuing from wherever its counter happened to be.

That explicit epoch argument is deprecated in torch 2.2.2, so a composed schedule emits the deprecation warning from inside the library, once per hand-over, on a line you never wrote. Current torch calls a private _update_lr(0) in the same branch instead, which does the same reset without the warning.

The reason to know this is diagnostic. A warning that names scheduler.step() and the epoch parameter, fired by a program that never passes one, is SequentialLR reaching a milestone, not a bug in your loop.

    def step(self):
        self.last_epoch += 1
        idx = bisect_right(self._milestones, self.last_epoch)
        scheduler = self._schedulers[idx]
        if idx > 0 and self._milestones[idx - 1] == self.last_epoch:
            scheduler.step(0)
        else:
            scheduler.step()

        self._last_lr = scheduler.get_last_lr()

The chapter's rule names the RNG as part of the state when bitwise resume matters. Here is what it weighs: torch.get_rng_state() returns a uint8 tensor of 5056 elements, the CPU generator's Mersenne Twister block, and it goes into the checkpoint dict like anything else.

The proof is worth running rather than reading. Twelve steps with dropout in the model and a randint drawing the batch, checkpointed at step six, restored into a fresh model, optimizer and scheduler, with the RNG set back. The six resumed losses equal the last six of the run that never stopped, as a list comparison, not as a curve that looks about right.

Equality rather than similarity is the bar because a bitwise resume is a claim you can check in one line. A curve that overlays plausibly is compatible with a scheduler that reset, a dropout mask that shifted, and an optimizer moment that came back empty.

import torch
from torch import nn
from torch.optim.lr_scheduler import CosineAnnealingLR


def build():
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Dropout(0.2), nn.Linear(32, 1))
    opt = torch.optim.AdamW(model.parameters(), lr=1e-2)
    return model, opt, CosineAnnealingLR(opt, T_max=12)


def step(model, opt, sched):
    idx = torch.randint(0, 256, (32,))              # this line reads the global RNG
    loss = nn.functional.mse_loss(model(data[idx]), target[idx])
    opt.zero_grad()
    loss.backward()
    opt.step()
    sched.step()
    return round(loss.item(), 6)


torch.manual_seed(1)
data = torch.randn(256, 8)
target = data.sum(dim=1, keepdim=True)

model, opt, sched = build()                         # the run that never stops
torch.manual_seed(7)
straight = [step(model, opt, sched) for _ in range(12)]

model, opt, sched = build()                         # the same run, killed at step 6
torch.manual_seed(7)
for _ in range(6):
    step(model, opt, sched)
torch.save(
    {
        "model": model.state_dict(),
        "opt": opt.state_dict(),
        "sched": sched.state_dict(),
        "step": 6,
        "rng": torch.get_rng_state(),
    },
    "ckpt.pt",
)

state = torch.load("ckpt.pt")
model, opt, sched = build()
model.load_state_dict(state["model"])
opt.load_state_dict(state["opt"])
sched.load_state_dict(state["sched"])
torch.set_rng_state(state["rng"])
print([step(model, opt, sched) for _ in range(6)] == straight[6:])   # True
print(state["rng"].dtype, state["rng"].shape)   # torch.uint8 torch.Size([5056])

Drop one piece at a time from that restore and the damage has a shape. Leave the RNG out and the very first resumed step is already wrong, 8.478 against 7.358, because a different batch got drawn and a different dropout mask fired.

Leave the scheduler out and the first two steps match exactly, then the third drifts. The optimizer's own state dict carries the current lr, so step seven runs at the correct 0.005 by accident. From step eight the fresh scheduler walks the cosine from its beginning: 0.004915 where the real run had 0.003706, and by step twelve 0.003147 against 0.00017, more than eighteen times the intended rate.

That delay is the part to remember. A resume bug that is invisible for two steps and then bends the curve is a bug people attribute to data, to hardware, to anything except the two lines that restored the wrong object.

def resume(restore_rng, restore_sched):
    state = torch.load("ckpt.pt")        # reloaded each time, for lesson one's reason
    model, opt, sched = build()
    model.load_state_dict(state["model"])
    opt.load_state_dict(state["opt"])
    if restore_sched:
        sched.load_state_dict(state["sched"])
    if restore_rng:
        torch.set_rng_state(state["rng"])
    return [step(model, opt, sched) for _ in range(6)]


print(straight[6:])
print(resume(True, True))
# [7.357611, 6.734547, 9.127593, 8.237515, 8.615776, 5.698226]  on both lines
print(resume(False, True))
# [8.477634, 4.781869, 6.414695, 6.638178, 9.703715, 5.7183]
print(resume(True, False))
# [7.357611, 6.734547, 9.101551, 8.163261, 8.458454, 5.567288]

Every piece so far has a state dict. The input pipeline does not, and chapter five's shuffling is the reason it matters here.

A DataLoader with shuffle=True draws its permutation from a generator, and you can pass your own. Save that generator's state at the end of an epoch and the next epoch's batch order is reproducible: the same four batches, in the same order, from a state you saved yourself.

Stopping mid-epoch is the case with no answer. Nothing records how many batches were consumed, so a resume from the same generator state gets a fresh permutation of the whole dataset rather than the batches you had left. Two batches into a four-batch epoch, the resume hands you all twelve samples again in a new order, and some rows get seen twice this epoch while others wait.

import torch
from torch.utils.data import DataLoader, TensorDataset

ds = TensorDataset(torch.arange(12.).unsqueeze(1))


def epoch(rng_state, stop_after):
    g = torch.Generator()
    g.set_state(rng_state)
    out = []
    for i, (b,) in enumerate(DataLoader(ds, batch_size=3, shuffle=True, generator=g)):
        if i == stop_after:
            break
        out.append([int(v) for v in b.flatten()])
    return out, g.get_state()


start = torch.Generator().manual_seed(0).get_state()
first, after_first = epoch(start, 4)
second, _ = epoch(after_first, 4)
print(first)    # [[5, 4, 0], [9, 7, 8], [1, 10, 3], [11, 6, 2]]
print(second)   # [[1, 11, 0], [3, 5, 4], [2, 7, 6], [9, 10, 8]]
print(epoch(after_first, 4)[0] == second)   # True: the order is recoverable

half, mid = epoch(start, 2)     # killed two batches into epoch one
print(half)                     # [[5, 4, 0], [9, 7, 8]]
print(epoch(mid, 4)[0])         # [[8, 6, 10], [0, 5, 11], [4, 3, 2], [1, 9, 7]]

Five things, four of which serialize and one of which does not. The table is the whole answer to what a checkpoint has to contain, and the last column is what this lesson measured rather than what it assumed.

One version difference belongs in the same breath, because it changes how the file comes back. On torch 2.2.2 torch.load defaults to weights_only=False; current torch defaults it to True, documented as restricting the unpickler to tensors, primitive types, dictionaries and anything you explicitly allow. A checkpoint shaped like the one above is tensors, ints and dicts, and it loads under weights_only=True on this machine when the flag is passed by hand. A checkpoint that pickled an optimizer object or an argparse namespace does not.

The TPU version of this exercise, kill the run, resume, and overlay the two curves, is the capstone in LAB·P4 and chapter twelve's mastery bar. What changes over there is the device and the bridge. What does not change is the list below.