DL S023 · Schedulers — Warmup, Cosine, One-Cycle, and WSD

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

The story · "just pick 3e-4" and the moment that stopped working

For a long time, "learning rate" was one number and one folk theorem. You picked 3e-4 (the "Karpathy constant" for Adam, semi-jokingly attributed to Karpathy but actually just a common default in the 2015-2017 era) and hoped. If loss went NaN, you halved it. If loss barely moved, you doubled it. That was the whole science. CIFAR ResNets, ImageNet ResNets, small LSTM language models — one constant learning rate, one number, worked more often than it did not.

Then transformers happened. In late 2017, the Attention Is All You Need team noticed that their model diverged spectacularly if they started training with any reasonable constant learning rate. The first few hundred updates were so large that the model landed in a region of the loss surface it never recovered from. Their fix was silly-simple, buried in section 5.3 of the paper: warm up. Start LR from ~0 for the first N steps, ramp linearly to your target, then decay proportional to 1/√step. Suddenly transformers trained. The linear-warmup-then-inverse-square-root-decay schedule became known as the "Noam schedule" (after Noam Shazeer, the paper's fifth author, who had introduced the idea in earlier Google internal work). Every LLM paper since GPT-2 uses some flavor of warmup-then-decay. Modern LLMs use cosine decay instead of inverse-square-root, but the ancestor is the Noam schedule.

Fast-forward to 2024. Two very different schedule innovations shook things up:

• Warmup-Stable-Decay (WSD) — from MiniCPM (Hu et al., April 2024). Instead of warmup → cosine → done, use warmup → long constant plateau at max LR → rapid decay in the last 10-20% of training. This schedule (a) beats cosine at Chinchilla-optimal ratios, (b) lets you train an "infinite" model where the decay phase is applied lazily whenever you decide to release a checkpoint, and (c) makes "continued pretraining" much cleaner because you can pause anywhere in the stable phase and later resume.
• Muon optimizer — Jordan et al. October 2024 — a new optimizer that replaces Adam for hidden-layer matrices in transformers by projecting each gradient onto its "orthogonalized" (whitened) direction using a Newton-Schulz iteration. It changes what "learning rate" even means and often converges 30-50% faster than AdamW at equal loss. Muon does not need warmup for the projected updates (they are already bounded by construction) but still uses it for the embedding and output layers which are trained with AdamW. Fascinatingly, the 2024 Kimi K2 tech report mentions Muon in their optimizer ablations.

Today: what a scheduler is, the four schedules that cover 99% of use cases (constant, warmup+cosine, one-cycle, WSD), how to wire them into a training loop correctly, and the ten most common bugs.

Prerequisites

• Session 017 — training loop skeleton.
• Session 018 — AMP wrapping does not change scheduler wiring, but you should know they compose.
• Session 022 — bad init compounds with bad early LR into instant NaN.

1 · A scheduler is a rule that mutates optimizer.param_groups[i]['lr'] over time

PyTorch schedulers are shockingly simple. They hold a reference to your optimizer and, on every .step(), mutate its param_groups[i]['lr'] according to some formula.

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=1000)
 
for step in range(1000):
    ...
    optimizer.step()
    scheduler.step()             # update LR for the NEXT iteration

That is it. scheduler.step() computes a new LR based on internal state and the schedule formula, then sets it in place on the optimizer's param groups. When optimizer.step() runs on the next iteration, it uses this updated LR.

Rule: always call scheduler.step() AFTER optimizer.step(). Reverse order used to be OK in old PyTorch versions; now it prints a warning and gives subtly wrong behavior (the very first optimizer step uses whatever the scheduler set on init, not what you expect).

2 · The four schedules you will actually use

2.1 Constant LR — still the right default for many tasks

If your model is small, your training is short, and your data is clean — pick a constant LR from a quick sweep (3e-4, 1e-3, 3e-3 for AdamW is usually the range) and stop thinking. Fine-tuning a small classification head on top of frozen embeddings? Constant is fine. Toy notebook? Constant is fine. The bugs that constant LR causes are the easiest to diagnose.

2.2 Warmup + linear decay ("BERT / OG transformer schedule")

Linear ramp UP for warmup_steps, then linear ramp DOWN to 0 over the rest. Every LLM pretrain from GPT-2 through LLaMA-2 uses some flavor of this. Simple, predictable, easy to explain in a paper's methods section.

from transformers import get_linear_schedule_with_warmup
scheduler = get_linear_schedule_with_warmup(
    optimizer, num_warmup_steps=2000, num_training_steps=100_000
)

Or in pure PyTorch:

def linear_warmup(step, warmup, total):
    if step < warmup:
        return step / warmup
    return max(0.0, (total - step) / (total - warmup))
 
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda s: linear_warmup(s, 2000, 100_000))

2.3 Warmup + cosine decay (modern LLM default)

Same warmup, but then decay along a cosine from lr_max to lr_min (usually lr_min = 0.1 · lr_max). Smoother than linear at the end of training; slightly better final loss in practice.

import math
 
def cosine_lr(step, warmup, total, lr_max, lr_min):
    if step < warmup:
        return lr_max * step / warmup
    progress = (step - warmup) / (total - warmup)
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * progress))

Used by GPT-3, PaLM, LLaMA, Mistral, Gemma, DeepSeek-V3 — this was the frontier default from 2020 through mid-2024. It is still the safe default for a first pretraining run.

The lr_min = 0.1 · lr_max is not arbitrary. Cosine to zero (lr_min = 0) leaves the model in a state where the last 5% of training does effectively nothing (LR is too small to matter), so you have wasted 5% of your compute. lr_min = 10% of max keeps the model learning throughout while still cooling.

2.4 One-Cycle (Leslie Smith, "super-convergence")

Ramp UP to lr_max, then DOWN below the initial LR, all inside one training run. Momentum moves OPPOSITE to LR (down when LR is up, up when LR is down). Empirically this trains CNNs 2-4× faster than constant LR on ImageNet (Smith 2018).

scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=1e-2, total_steps=len(loader) * epochs
)

Great for image classification with limited compute budget. Less used for LLMs (cosine is the default there — one-cycle's aggressive ramp-down does not help when your total training is compute-constrained rather than time-constrained).

2.5 Warmup-Stable-Decay (WSD, the 2024 challenger)

The 2024 MiniCPM paper (Hu et al.) introduced Warmup-Stable-Decay:

where f is an exponential or 1/√ decay to ~10% of max over the last 10-20% of steps.

Three surprising properties:

1. Beats cosine at Chinchilla-optimal training. MiniCPM's 1.2B model trained on 1T tokens with WSD outperformed the same architecture trained on 1T tokens with cosine, using the same peak LR. The stable phase is where most learning happens; the decay phase is a "cool-down" that helps the model settle.
2. "Infinite" training / lazy decay. During the stable phase, no decay is happening. You can pause anywhere in the stable phase, snapshot the model, apply the decay phase lazily to that snapshot, and get a shippable checkpoint. This lets one very long training run produce many "final" models at different token budgets — an enormous logistical win.
3. Clean continued pretraining. If you have an existing WSD-trained model and want to train longer, you extend the stable phase and reapply the decay — much cleaner than resuming a cosine schedule (which has no obvious extend semantics).

WSD implementation (using LambdaLR):

def wsd_lr(step, warmup, stable_end, total, min_ratio=0.1):
    if step < warmup:
        return step / warmup
    if step < stable_end:
        return 1.0
    # Exponential decay from 1.0 to min_ratio over (total - stable_end) steps
    progress = (step - stable_end) / (total - stable_end)
    return math.exp(math.log(min_ratio) * progress)
 
# Example: 1B token training, 1% warmup, decay in last 10%
scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer, lambda s: wsd_lr(s, warmup=1_000, stable_end=90_000, total=100_000)
)

WSD has been adopted by MiniCPM (obviously), some parts of the 2024-2025 OLMo 2 recipe, and various academic reproductions. It is not (yet) the default for the largest frontier models (Llama 3, DeepSeek-V3 both used cosine) but is gaining share fast.

3 · Why warmup? — the Adam variance argument, expanded

At step 1, Adam has almost no history in its momentum buffers (m and v). Its variance estimate v is tiny, so the update -lr · m / (√v + eps) is enormous (dividing by a small √v). Combined with random-init weights that produce large loss gradients, this launches parameters into a bad region of the loss landscape that they never recover from.

Concretely, at step 1 of Adam:

So Adam's first update is essentially η · sign(g) — a step of size exactly η in the direction of the gradient. At η = 6e-4 and 175B parameters, that is 105M worth of unit-magnitude parameter changes in one step. If any of them land badly, the model is unrecoverable.

Warmup fixes this by making the first update tiny (η ≈ 0 at step 1), then slowly ramping up so that by the time η is at its target value, v has accumulated enough history that the raw sign() behavior is smoothed out.

Also relevant: BN/LN running-stats need a few steps to stabilize. Warmup gives them room too.

Rule of thumb for LLM training: warmup_steps ≈ 0.5%–5% of total steps. GPT-3 used 375M tokens of warmup out of 300B — about 0.1%. BERT used 10K steps out of 1M — 1%. Llama 3 used 8K steps out of ~1M — under 1%. Pick your battles: too little warmup and Adam blows up; too much and you have wasted early compute at LRs so small nothing is happening.

Related: the RAdam alternative. Liu et al. 2019 proposed "Rectified Adam," which computes an analytic estimate of when the variance is trustworthy and adaptively rectifies the update magnitude — making warmup unnecessary. It works but never really caught on because "add warmup" is one line and "switch optimizers" is a bigger config change.

4 · The full transformer-style loop

Read carefully — the step cadence is where people mess up.

import torch, math
 
model = ...
optimizer = torch.optim.AdamW(model.parameters(), lr=6e-4, weight_decay=0.1,
                              betas=(0.9, 0.95))
 
total_steps  = len(train_loader) * epochs
warmup_steps = int(0.01 * total_steps)     # 1% warmup
 
def lr_at(step):
    if step < warmup_steps:
        return step / warmup_steps         # ramp 0 → 1
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return max(0.1, 0.5 * (1 + math.cos(math.pi * progress)))   # cosine to 10%
 
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_at)
 
global_step = 0
for epoch in range(epochs):
    for xb, yb in train_loader:
        optimizer.zero_grad(set_to_none=True)
        loss = criterion(model(xb), yb)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()                    # PER STEP, not per epoch
        global_step += 1
        writer.add_scalar("train/lr", scheduler.get_last_lr()[0], global_step)

Log the LR to TensorBoard (or W&B, or a file). Every serious training run. It is a two-line addition that catches "why isn't my LR changing?" in seconds. Modern practice: also log grad norm, weight norm, and per-parameter-group LR (if you have different LRs for different groups).

5 · The step-cadence trap

Some schedulers are meant to be stepped per batch; others per epoch. Get this wrong and your effective schedule spans the wrong horizon.

Rule: batch-cadence for continuous schedules (cosine, linear, WSD, one-cycle), epoch-cadence for step-based ones (StepLR, MultiStepLR, ReduceLROnPlateau).

If you are using gradient accumulation (Session 024), step the scheduler once per optimizer step, not once per micro-batch. That is: after every accum_steps micro-batches you do one optimizer.step() + one scheduler.step().

6 · Resume-from-checkpoint MUST save the scheduler

Session 017 §4 flagged this: a checkpoint should include scheduler.state_dict(). Without it, resume starts the schedule from scratch (LR back at warmup start!), and the next few thousand steps use a wildly wrong LR.

Every scheduler has a state_dict — always save it, always load it. This is one of the top-3 causes of failed training runs in production, tied with silent data corruption and OOM at hour 47.

7 · ReduceLROnPlateau — the "poor man's" scheduler

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode="min", factor=0.5, patience=3
)
 
# At end of each epoch:
scheduler.step(val_loss)

If val_loss has not improved for patience epochs, LR is multiplied by factor. Used in fine-tuning, small experiments, when you do not know total_steps up front. Not what LLM pretraining uses — but a lifesaver on smaller research projects where you would rather the scheduler be data-driven than schedule-driven.

Watch out for two footguns: (1) ReduceLROnPlateau doesn't have a get_last_lr() method in older PyTorch, so you have to inspect optimizer.param_groups[0]['lr'] directly to log; (2) it uses raw comparisons of val_loss, so if your val loss is noisy, patience of 3 is too aggressive — bump to 5-10 or add a threshold argument.

8 · The 2024-2025 optimizer landscape — where scheduling gets interesting

8.1 AdamW is still king for language modeling

Every frontier LLM in 2024-2025 uses AdamW (or a fused implementation of it) with betas=(0.9, 0.95) and weight_decay=0.1. This has not changed in 5 years, and probably won't change in 2026 either — AdamW's combination of "sensible defaults work" and "billions of GPU-hours of tuning wisdom baked in" is very hard to beat.

8.2 Lion — Google Brain's 2023 challenger

Symbolic Discovery of Optimization Algorithms (Chen et al. 2023) — Google's Lion optimizer, discovered by a large program-search over optimizer space, uses sign() of the momentum for the update (no per-parameter variance estimate) and is ~15% more memory-efficient than AdamW at similar quality. Used in some Google-adjacent projects, not widely adopted. Same schedule (warmup + cosine or WSD) works.

8.3 Muon — the 2024 hidden-layer whitening idea

Muon: MomentUm Orthogonalized by Newton-Schulz (Jordan et al. 2024) — for weight matrices only, replace the raw gradient update with the "orthogonalized" gradient (U @ V^T from the SVD g = U Σ V^T, computed cheaply via 5 iterations of Newton-Schulz). Empirically 30-50% faster convergence than AdamW on hidden-layer matrices in transformers. Embeddings and output projections still use AdamW.

Interesting scheduling implication: Muon's updates are bounded by construction (they are orthogonal, so bounded singular values), so it needs much less warmup than AdamW — some Muon runs use no warmup at all for the Muon-optimized params, only warmup on the AdamW-optimized params.

Public reports:

• The Kimi K2 tech report (2025) mentions Muon in their ablations and adopts it for hidden layers.
• Community reproductions on nanoGPT + Muon show Muon reaching GPT-2 loss in 3× fewer steps than baseline.

Muon is still early — it will either become standard by 2027 or fade like Lion did. Watch this space.

8.4 The lesson: pick your schedule based on your optimizer

AdamW → warmup + cosine or WSD. Lion → same. SGD with momentum → one-cycle. Muon → minimal warmup for Muon params, standard warmup for AdamW params. Read the optimizer's paper's schedule section — they've done the tuning.

Further reading

• Hu et al. 2024, MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies (WSD).
• Loshchilov & Hutter 2016, SGDR: Stochastic Gradient Descent with Warm Restarts (cosine).
• Smith 2018, A disciplined approach to neural network hyper-parameters (one-cycle).
• Liu et al. 2019, On the Variance of the Adaptive Learning Rate and Beyond (RAdam, warmup theory).
• Jordan et al. 2024, Muon: MomentUm Orthogonalized by Newton-Schulz.

9 · War stories

10 · Diagram — the five common shapes

Every plot: y-axis LR, x-axis step. Once you see the shapes, papers become much easier to translate — "we use WSD with 1% warmup, 89% stable, 10% decay" is one line to read.

11 · Try it yourself

12 · Recap

A scheduler is a rule that mutates LR over training. Constant is the trivial baseline. Warmup fixes the Adam-variance-at-step-1 problem for large models. Cosine to 0.1 · lr_max is the safe modern default. One-cycle wins for CNN classification with limited compute. WSD is the 2024 challenger: warmup → long constant plateau → rapid decay, beats cosine at Chinchilla-optimal training and enables "infinite" pretraining. Save scheduler.state_dict() in every checkpoint. Step the scheduler after the optimizer, once per optimizer step, at the same cadence you configured T_max for.

🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §3 (warmup derivation) and §2.5 (WSD) in 24 hours. Revisit the full session on day 7 with focus on §8 (2024 optimizer landscape).

Next session (S024): you have got init, norms, dropout, and schedules. What if a gradient is 100× too big and blows through your careful schedule anyway? Session 024 is gradient clipping and gradient accumulation — clipping to survive occasional huge updates (a must for LLM training), and accumulation to fit big effective batches on small GPUs (a must for laptop training and the trick that makes DDP practical at scale).

Sticky note (keep on your desk):

• Warmup + cosine is the LLM default; 0.5%-5% of total steps for warmup, lr_min = 0.1 · lr_max.
• WSD if you want to extend training or ship checkpoints lazily; 80-90% stable phase, 10-20% decay.
• scheduler.step() AFTER optimizer.step(), per BATCH for continuous schedules.
• Always save scheduler.state_dict() in checkpoints.

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