Schedulers — Warmup, Cosine, One-Cycle, and WSD

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 23 / 80 · Module M04 · ~3 hours

0 · A schedule needs an explicit clock

A learning-rate schedule is a time-indexed policy. The central engineering question is what counts as time: consumed examples/tokens, attempted steps, accepted optimizer updates or completed epochs. A correct formula on the wrong clock is still a wrong experiment.

Warmup, cosine, linear decay, one-cycle and warmup-stable-decay are useful shapes—not universal rankings or guarantees against divergence. This lesson derives their endpoints, executes schedule/resume controls and compares four schedules on the same small synthetic task. No GPU, pretrained model or dataset download is needed.

Prerequisites

Know the training loop, precision/scaler update boundary and initialization assumptions. PyTorch API examples target 2.9. All losses use natural logarithms.

1 · A scheduler changes optimizer parameter-group learning rates

An optimizer reads param_groups[i]['lr'] when it updates parameters. A scheduler mutates those values according to a policy and saved state. For ordinary update-clock schedulers, call optimizer step first, then scheduler step to prepare the next update. The scheduler's construction can already set the initial LR; for LambdaLR, multiplier f(0) takes effect before the first update.

last_epoch is a historical name for the scheduler's counter, not proof it must count epochs. base_lrs, group ordering and the defining callable/config matter too. For logging, record used_lr before optimizer.step and distinguish it from the next LR returned after scheduler.step. Calling a scheduler before the optimizer can skip the intended first value.

2 · Five useful shapes and explicit endpoint conventions

Let t count completed optimizer updates, with planned horizon T and warmup endpoint W. The next optimizer update uses the LR at t; construction uses t=0, so a zero-start warmup intentionally makes the first update zero-LR. This is still an optimizer call: AdamW can advance its moment estimates and step counter even though parameters do not move. Here the clock counts successful calls, not only calls with a nonzero parameter delta. Another valid convention is first-update LR $\eta_{max}/W$, but it must be declared and tested rather than mixed into the same formula. After T, this lesson's custom functions clamp, not rebound.

2.1 Constant LR

Constant is a baseline, especially for short/local runs. Select it by a bounded sweep, not faith in 3e-4. It still needs optimizer state and consumed-step/data state for continuation; 'no schedule' does not mean 'nothing to restore'.

2.2 Warmup plus linear decay

For $0\le W<T$,

When W=0, skip the first branch. This is a piecewise formula, not a min expression that should divide by zero when warmup is absent. A nonzero floor can be included in the second branch. Hugging Face has a get_linear_schedule_with_warmup helper, but the complete pure-PyTorch implementation below avoids adding that dependency.

The original Transformer's warmup/inverse-square-root schedule is not linear decay: it used $d_{model}^{-1/2}\min(t^{-1/2},tW^{-3/2})$ for positive t. BERT-style linear schedules and later cosine recipes must not be assigned indiscriminately to every LM family.

2.3 Warmup plus cosine decay

After warmup define $u=(t-W)/(T-W)$. Then

At u=0, midpoint .5 and endpoint 1, the values are max, (max+min)/2 and min. max(min_ratio, .5*(1+cos(...))) is not the same function: it clips a zero-floor cosine into a flat tail. For max=6e-4 and min=6e-5, the proper midpoint is 3.3e-4, not 3e-4.

Zero-floor cosine is not automatically five percent wasted compute; near-zero LR can serve final annealing, while an overly large floor can hurt final quality. Choose endpoints by the task and budget. Raw CosineAnnealingLR does not enforce a terminal clamp; stepping beyond its intended T_max can increase LR again even though it is not the separate warm-restarts scheduler. Stop at the planned horizon or implement an explicit extension policy.

2.4 One-cycle

OneCycleLR raises LR from an initial value to a peak then lowers it below the initial value. With default options, initial LR is max_lr/25, final minimum is that initial value divided by 10,000, pct_start=.3, and momentum cycles inversely when enabled. PyTorch's default two-phase policy differs from the paper-like three_phase=True option. The API is not chainable in the same way as arbitrary schedulers.

Count actual optimizer updates in total_steps, not microbatches under accumulation. Match max LR, momentum policy and phase lengths in comparisons. Some studies find rapid convergence, but no fixed 2–4× ImageNet speedup or universal CNN-versus-LM ranking follows.

2.5 Warmup-stable-decay (WSD)

WSD explicitly separates warmup, a constant high-LR interval and a final decreasing interval. This lesson's concrete tail is exponential: if decay starts at S and minimum ratio is r,

It reaches 1 at S and r at T, and clamps afterward. A linear or cosine tail is also a WSD variant; the name does not uniquely identify the function. MiniCPM's paper defines a general decreasing f and studies explicit phase budgets. Its small-model experiments found a roughly 10%-token decay sufficient in the tested setting; that is not a mandatory 80–90% stable rule for every model.

The useful planning feature is branching: retain a full checkpoint from the stable phase, continue one branch at high LR and decay another for evaluation. Decay still costs training data and compute; it is not a postprocessing transform on fixed weights. Extending from the already cooled checkpoint is a different experiment from returning to the stable checkpoint. The paper's 0.036B experiments and inferred scaling envelope do not establish an invented 1.2B/1T universal win over cosine.

import math
 
def ratio(step,warmup,total,kind="cosine",minimum=.1,stable_end=None):
    if not isinstance(step,int) or step < 0: raise ValueError("nonnegative integer clock")
    if not 0 <= warmup < total: raise ValueError("need 0 <= warmup < total")
    if not 0 <= minimum <= 1: raise ValueError("floor ratio")
    if kind not in {"constant","linear","cosine","wsd"}: raise ValueError("kind")
    if kind == "wsd" and not (stable_end is not None and warmup <= stable_end < total and minimum > 0):
        raise ValueError("WSD needs positive floor and a nonempty tail")
    if warmup and step < warmup: return step/warmup
    if kind == "constant": return 1.
    t=min(step,total)
    if kind == "wsd":
        if t <= stable_end: return 1.
        return math.exp(math.log(minimum)*(t-stable_end)/(total-stable_end))
    progress=(t-warmup)/(total-warmup)
    if kind == "linear": return minimum+(1-minimum)*(1-progress)
    return minimum+.5*(1-minimum)*(1+math.cos(math.pi*progress))
 
assert ratio(0,0,10) == 1.
assert ratio(0,2,10) == 0. and ratio(2,2,10) == 1.
assert abs(ratio(6,2,10)-.55) < 1e-12
assert abs(ratio(10,2,10)-.1) < 1e-12
assert ratio(100,2,10) == ratio(10,2,10)
assert abs(ratio(9,2,10,"wsd",.1,8)-math.sqrt(.1)) < 1e-12
assert ratio(8,2,10,"wsd",.1,8) == 1.
proper=6e-4*ratio(6,2,10)
wrong=6e-4*max(.1,.5*(1+math.cos(math.pi*.5)))
assert abs(proper-3.3e-4) < 1e-15 and abs(proper-wrong) > 2e-5
for bad in [(0,10,10),(0,-1,10)]:
    try: ratio(*bad)
    except ValueError: pass
    else: raise AssertionError("invalid horizon accepted")
print("piecewise boundaries and clipped-cosine negative control passed")

3 · Why warmup? Derive Adam before telling a story

At the first step, ordinary Adam's bias-corrected moments are $\widehat m_1=g_1$ and $\widehat v_1=g_1^2$. Ignoring decay,

Each coordinate's magnitude is at most eta, not arbitrarily huge merely because uncorrected v started small. For gradients well above epsilon it is approximately an eta-sized sign update. Many such coordinate changes can still be large relative to initial weights or collectively large: for P comparable coordinates the L2 update is approximately $\eta\sqrt P$, while the L1 update is $\eta P$. Summing 175B coordinate magnitudes is not a direct divergence criterion.

Short histories and adaptive scaling, initialization, residual/attention sensitivity, batch size and changing gradient distributions can make warmup useful. RAdam analyzes variance in adaptive learning-rate estimates under assumptions; it does not prove warmup belongs exclusively to Adam or that RAdam makes it unnecessary everywhere. Goyal et al.'s large-batch SGD study explicitly uses warmup, refuting that exclusive claim. LayerNorm has no running statistics to warm up; BatchNorm may have them, but it is not a general transformer explanation.

A warmup budget in updates, tokens or epochs can be legitimate if its conversion is explicit. Changing batch size while retaining 'one epoch warmup' changes update count. The quantity $1/(1-\beta_2)$ is an EMA timescale, not a universal formula for necessary warmup independent of architecture/data.

import torch
from torch import nn
 
p=nn.Parameter(torch.zeros(4,dtype=torch.float64))
g=torch.tensor([1.,-.1,1e-10,0.],dtype=torch.float64)
opt=torch.optim.Adam([p],lr=.001,eps=1e-8)
p.grad=g.clone();opt.step()
expected=-.001*g/(g.abs()+1e-8)
assert torch.allclose(p,expected,atol=1e-15,rtol=1e-12)
assert p.abs().max() <= .001
# Wrong uncorrected-moment story predicts sqrt(.001) denominator factor.
wrong=-.001*(.1*g)/(torch.sqrt(.001*g*g)+1e-8)
assert not torch.allclose(p,wrong)
print("bias-corrected Adam first update",p.detach().tolist())

4 · Full update-clock loop with logging and resume

The complete CPU program uses the tested cosine function, AdamW, clipping, validation and a resume checkpoint. It saves the defining schedule configuration because a plain lambda's code is not serialized by LambdaLR. Construct model, optimizer and scheduler before loading; restore scheduler state and then optimizer state so construction does not overwrite the saved LR. Identical logs/weights and held-out loss are tested at the interrupted boundary. This bounded example recreates fixed synthetic data from its seed and has no dropout, shuffle or random augmentation after initialization; it is not a substitute for saving RNG/sampler state in a stochastic training job. The held-out examples are never used by the optimizer.

import copy
import torch
from torch import nn
import torch.nn.functional as F
 
def scheduled_run(limit=30,resume=None,reset_schedule=False):
    torch.manual_seed(23)
    model=nn.Sequential(nn.Linear(4,12),nn.Tanh(),nn.Linear(12,2))
    x=torch.randn(20,4);y=(x[:,0] > 0).long()
    vx=torch.randn(60,4);vy=(vx[:,0] > 0).long()
    opt=torch.optim.AdamW(model.parameters(),lr=.03,weight_decay=.01)
    config={"warmup":3,"total":30,"minimum":.1}
    sched=torch.optim.lr_scheduler.LambdaLR(opt,lambda t:ratio(t,**config))
    start=0;log=[]
    if resume:
        assert resume["config"] == config
        model.load_state_dict(resume["model"])
        if not reset_schedule: sched.load_state_dict(resume["sched"])
        opt.load_state_dict(resume["optim"])
        start=resume["updates"];log=copy.deepcopy(resume["log"])
    for update in range(start,limit):
        model.train();opt.zero_grad(set_to_none=True)
        loss=F.cross_entropy(model(x),y);loss.backward()
        norm=torch.nn.utils.clip_grad_norm_(model.parameters(),1.,error_if_nonfinite=True)
        used=opt.param_groups[0]["lr"]
        opt.step();sched.step()
        log.append((update,used,loss.item(),norm.item(),sched.get_last_lr()[0]))
    model.eval()
    with torch.no_grad(): metric=F.cross_entropy(model(vx),vy).item()
    return copy.deepcopy(dict(model=model.state_dict(),optim=opt.state_dict(),
                              sched=sched.state_dict(),updates=limit,log=log,
                              config=config,metric=metric))
 
full=scheduled_run();part=scheduled_run(11);resumed=scheduled_run(resume=part)
assert full["log"] == resumed["log"]
assert full["metric"] == resumed["metric"]
assert all(torch.equal(full["model"][k],resumed["model"][k]) for k in full["model"])
assert full["log"][0][1] == 0. and abs(full["log"][3][1]-.03) < 1e-12
assert scheduled_run(resume=part,reset_schedule=True)["log"] != full["log"]
print("scheduler continuation and reset-negative-control passed",full["metric"])

5 · Cadence, accumulation and skipped updates

StepLR(step_size=E) reduces every E scheduler calls; using epoch calls is a convention. Cosine's T_max likewise counts calls and can be epoch- or update-based. OneCycleLR is documented per training batch; under accumulation, configure it for optimizer updates rather than treating every microbatch as a new optimizer batch. ReduceLROnPlateau consumes the named validation metric at evaluation events.

If 17 microbatches are grouped in fours and the tail is flushed, one epoch has five optimizer updates, not four and not seventeen. Unequal microbatch token counts change the loss denominator but not that count. Under FP16, GradScaler may skip an update; an accepted-update schedule should skip too. A token-clock schedule may deliberately continue because tokens were consumed—choose and document the policy.

For a single optimizer under ordinary GradScaler growth/backoff, comparing scale before and after update() can indicate a skipped update when the scale decreased. This is not a general event API for multiple optimizers or manually overridden scales. optimizer.step() commonly returns None on success, so None is not evidence it was skipped. The precision lesson supplies the complete reference and CPU skip control.

6 · Resume needs policy as well as scheduler state

Save model, optimizer, scheduler, scaler, data/RNG position and update/token/epoch counters. Record schedule type, warmup/tail boundaries, minimum, base/group rates and the code/version that defines a lambda. Loading scheduler.state_dict() alone does not restore the optimizer's moment buffers, current group LR or sampler. Restore the same group ordering.

Changing the remaining budget is not ordinary resume: a new cosine horizon bends the future curve, whereas continuing a stable WSD branch can preserve its already-used values. Store that intervention as a new experiment. Never relabel an already-decayed checkpoint as if it came from an uncooled stable branch.

7 · ReduceLROnPlateau — a metric-driven policy

Choose mode to match a metric, threshold and relative/absolute threshold mode to define meaningful improvement, patience, cooldown, factor, min_lr and eps. In PyTorch 2.9 the reduction occurs when bad-event count exceeds patience. With patience zero, the first event establishes a best baseline and the next non-improving event can reduce LR. Save best and bad-event/cooldown state.

Do not feed noisy training minibatch losses while describing validation patience. A metric-driven schedule is valid when the horizon is uncertain; it is not a 'poor man's' algorithm or intrinsically excluded from any model type. Too frequent reductions can react to noise; too much patience can miss useful adaptation.

import torch
from torch import nn
p=nn.Parameter(torch.tensor(1.));o=torch.optim.SGD([p],lr=1.)
s=torch.optim.lr_scheduler.ReduceLROnPlateau(o,factor=.5,patience=1,threshold=0.)
actual=[]
for metric in [1.,1.1,1.2,.9,1.,1.1]:
    s.step(metric);actual.append(o.param_groups[0]["lr"])
assert actual == [1.,1.,.5,.5,.5,.25]
# Built-in cosine does not clamp forever after its intended endpoint.
o=torch.optim.SGD([p],lr=1.);s=torch.optim.lr_scheduler.CosineAnnealingLR(o,T_max=2)
lrs=[o.param_groups[0]["lr"]]
for _ in range(4):
    p.grad=torch.zeros_like(p);o.step();s.step();lrs.append(o.param_groups[0]["lr"])
assert abs(lrs[2]) < 1e-12 and lrs[4] > .99
# OneCycle: initial and final USED rates, and the inverse momentum policy.
o=torch.optim.SGD([p],lr=.01,momentum=.9)
s=torch.optim.lr_scheduler.OneCycleLR(o,max_lr=.1,total_steps=20)
used=[]
for _ in range(20):
    used.append((o.param_groups[0]["lr"],o.param_groups[0]["momentum"]))
    p.grad=torch.zeros_like(p);o.step();s.step()
assert abs(used[0][0]-.004) < 1e-12
assert abs(used[-1][0]-4e-7) < 1e-12
peak=max(range(20),key=lambda i:used[i][0])
assert used[peak][1] < used[0][1]
print("plateau",actual,"cosine overrun",lrs,"one-cycle endpoints",used[0],used[-1])

8 · Optimizers change the scale the schedule acts on

8.1 AdamW

AdamW is a common adaptive baseline, not every frontier model's optimizer with identical betas/decay. Its decoupled weight decay multiplies weights separately from the adaptive gradient term. A schedule changes both gradient-step magnitude and the cumulative decay product under the usual implementation; adding L2 to ordinary Adam is not generally the same operation.

8.2 Lion

Lion uses a sign of a blend of the previous moment and current gradient for the parameter update, then updates the stored moment with a second beta. It does not keep Adam's second-moment buffer. That saves one parameter-sized state tensor; a claim like '15% less total memory' needs a stated denominator including parameters, gradients, precision copies and activations. Its recommended LR/decay scales differ from AdamW's; borrowing a schedule shape does not make the same numerical LR appropriate.

8.3 Muon and MuonClip

Muon applies an approximate matrix orthogonalization to a momentum-based update for appropriate hidden matrix parameters. The ideal SVD operation replaces singular values with ones in the corresponding semi-orthogonal factor; practical Newton–Schulz approximations, shape scaling and momentum conventions matter. Vectors and input/output layers commonly use another optimizer such as AdamW. This is not ordinary raw-gradient whitening or proof that warmup is unnecessary.

The old Muon citation 2410.14969 and Kimi K2 link 2504.12345 did not establish the claimed papers. The actual Kimi K2 report is 2507.20534 (2025) and describes MuonClip, adding QK clipping to address instability. That directly undermines the claim that bounded matrix-update singular values make the whole training system immune to instability. The author's Muon article reports particular tuned comparisons, not a universal 30–50% gain or forecast that it must dominate by 2027.

8.4 Match optimizer groups and schedule, then test

Choose LR magnitudes, decay and warmup per optimizer/parameter role. Momentum SGD can use constant, step, cosine or one-cycle; AdamW can use several shapes; Muon does not imply a zero-warmup rule. Record every group's LR. Distinguish convergence in updates, examples/tokens, FLOPs and wall time; 'fewer steps' is not necessarily faster or cheaper.

9 · Failure cases

• Diverges early without warmup: compare initialization, objective, precision and actual update/weight ratio. Adam's first bias-corrected coordinate is bounded by LR; do not blame an uncorrected denominator alone.
• Cadence wrong: a 100-call cosine stepped per batch reaches its endpoint after 100 batches, not 100 epochs, and may rise again if you keep stepping. Plot actual used rates.
• Resume restarts the curve: the negative control reproduces the error by omitting scheduler restoration. Restoring the correct counter without optimizer LR/state is also incomplete.
• Tiny final rate: inspect actual improvement and budget. A .1 floor is a hyperparameter, not a universal repair; a zero endpoint is not automatically wasted work.
• WSD tail too long or too short: compare the explicit phase allocation against a matched budget; the MiniCPM finding is study-bound, not a rule that all learning happens on the plateau.

11 · Try it yourself — four schedules, one local task

This synthetic comparison fixes data and initial model, uses AdamW throughout, logs LR and gradient norm, and reports clean held-out loss after an equal number of updates. The rates are deliberately specified per policy, so results reflect the whole configuration, not a pure shape-only causal effect. For a shape-only ablation sweep peak rate fairly and match other knobs. No winner or U-shaped result is asserted in advance.

import time
from pathlib import Path
import torch
from torch import nn
import torch.nn.functional as F
 
torch.manual_seed(31);x=torch.randn(64,4);y=(x[:,0]*x[:,1]>0).long()
vx=torch.randn(128,4);vy=(vx[:,0]*vx[:,1]>0).long()
curves={};grad_curves={};outcomes=[]
for name in ["constant","cosine","one-cycle","wsd"]:
    torch.manual_seed(4)
    model=nn.Sequential(nn.Linear(4,24),nn.Tanh(),nn.Linear(24,2))
    opt=torch.optim.AdamW(model.parameters(),lr=.01,weight_decay=.01)
    if name=="one-cycle":
        sched=torch.optim.lr_scheduler.OneCycleLR(opt,max_lr=.03,total_steps=80,cycle_momentum=False)
    else:
        sched=torch.optim.lr_scheduler.LambdaLR(opt,lambda t,n=name:ratio(t,0 if n=="constant" else 4,80,n,.1,72))
    lr_curve=[];grad_curve=[];start=time.perf_counter()
    for step in range(80):
        opt.zero_grad();loss=F.cross_entropy(model(x),y);loss.backward()
        grad=torch.nn.utils.clip_grad_norm_(model.parameters(),10.,error_if_nonfinite=True)
        lr_curve.append(opt.param_groups[0]["lr"]);grad_curve.append(grad.item())
        opt.step();sched.step()
    model.eval()
    with torch.no_grad(): val=F.cross_entropy(model(vx),vy).item()
    assert torch.isfinite(torch.tensor(val))
    outcomes.append((name,val,time.perf_counter()-start,grad.item()));curves[name]=lr_curve
    grad_curves[name]=grad_curve
    assert len(lr_curve)==len(grad_curve)==80
parts=['<svg xmlns="http://www.w3.org/2000/svg" width="720" height="300"><rect width="100%" height="100%" fill="white"/><text x="20" y="20">Actual learning rate used per optimizer update</text>']
parts.append('<path d="M30 45 V260 H510" fill="none" stroke="gray"/><text x="170" y="293">accepted optimizer update index</text>')
for i in [0,20,40,60,79]:
    parts.append(f'<text x="{25+i*6}" y="278" font-size="12">{i}</text>')
for val in [0,.01,.02,.03]:
    parts.append(f'<text x="0" y="{264-val/.03*210}" font-size="10">{val:.2f}</text>')
for j,(name,values) in enumerate(curves.items()):
    color=["#2563eb","#059669","#dc2626","#9333ea"][j]
    points=' '.join(f'{30+i*6},{260-v/.03*210}' for i,v in enumerate(values))
    parts.append(f'<polyline points="{points}" fill="none" stroke="{color}"/><text x="530" y="{60+j*30}" fill="{color}">{name}</text>')
parts.append('</svg>');Path('batch6-dl-training-schedules.svg').write_text(''.join(parts))
print("schedule, validation CE, CPU seconds, last preclip norm",outcomes)
# grad_curves[name][t] and curves[name][t] refer to the same update.
print("per-update preclip gradient norms",grad_curves)

For the original MNIST/CIFAR protocol, use an already authorized local dataset with fixed splits and repeated seeds, measure synchronized GPU wall time only if hardware is available, and report target-accuracy crossing as well as terminal loss. This CPU task checks the wiring and comparison discipline, not those datasets' outcomes.

12 · Recap and worked recall answers

Q1 — Why can Adam benefit from warmup? The first corrected update is -eta*g/(abs(g)+eps), not an unbounded reciprocal of tiny uncorrected v. Warmup controls early relative updates amid unstable estimates/representations; architecture and data matter. SGD can benefit too, and LN has no running statistics.

Q2 — Cosine versus one-cycle? Cosine lowers a rate over a named horizon; one-cycle first rises then falls below its initial rate, optionally cycling momentum inversely. Its default two-phase and paper-like three-phase settings differ. Which wins is an experiment, not an architecture rule.

Q3 — Batch or epoch? The declared clock decides. Convert horizon/warmup to the same units, count flushed accumulation tails, and decide how skipped AMP updates affect the clock. Log the actually used rate.

Q4 — What restores the schedule? Scheduler state plus its defining config/callable, optimizer state/group LR, counters, scaler and sampling/RNG state. The exact resumed program matches uninterrupted logs; resetting only the scheduler changes future rates.

Q5 — Two WSD planning properties? A high-LR stable phase can be extended without changing its past rates, and full stable checkpoints can branch into separately budgeted decay runs at several horizons. Those tails consume real training, and an already-decayed checkpoint is not the same branch origin.

Q6 — 3B model, 500B tokens, 512 H100s? If the budget is fixed and a validated cosine recipe exists, use it as a baseline. If continuation/multiple budgets matter, consider WSD while saving full stable checkpoints and reserving tail compute. Convert tokens to updates using actual supervised tokens per effective batch; GPU count alone does not determine schedule length. Compare pilot configurations rather than claim one mandatory answer or provision that hardware.

Sources and execution boundary

PyTorch scheduler source v2.9.0, optimizer docs 2.9, MiniCPM full paper v3, RAdam, large-batch SGD, Lion implementation, Muon author article, and Kimi K2 support the bounded claims. CPU formula, scheduler, resume and local comparison programs are executed; no LLM-scale schedule ranking, GPU benchmark or deployment is reproduced.

Optional pre-read resources

These inherited resource links are optional, not the primary evidence for this review. Video identity, duration and current availability are not certified.

• Learning Rate Schedulers in PyTorch
• Warmup and cosine schedules — LLM training practice
• Learning rate scheduling
• MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies

Previous: ← DL S022 · Weight Init · Next: DL S024 · Gradient Clipping + Accumulation →