Order the phases of a warmup-stable-decay pretraining run

Order steps·Easy·4.0 · 0·~1 min·Asked atGroq Notion

Attempt it

1Linear warmup: learning rate climbs from near zero while the batch size ramp completes
2Final checkpoint: the fully decayed model is evaluated and handed to post-training
3Anneal trigger: the data mixture shifts toward scarce high quality sources (textbooks, curated code, exam style data)
4Stable phase: long plateau at constant peak learning rate over the bulk of the token budget
5Decay phase: learning rate drops rapidly toward zero while the model digests the premium mix

**WSD** stands for warmup, stable, decay, and the order is in the name. First comes **warmup**. The LR climbs linearly from near zero over the first few thousand steps. Batch size is often ramped up in this same window. Warmup keeps update sizes sane while the weights move out of the chaotic region near initialization. Second comes the **stable phase**. The LR sits at its constant peak value over the bulk of the token budget. This is where most of the learning happens. Third, the **anneal trigger** fires. The data mixture shifts toward scarce high-quality sources: curated textbooks, vetted code, exam-style content. The shift is timed to the start of decay, not earlier, because the premium mix is most valuable when the model is consolidating. Fourth, the **decay phase**. LR drops rapidly toward zero while the model digests the premium mix. Fifth, the **final checkpoint** is evaluated and handed to post-training. The order is non-negotiable; rearranging it breaks the recipe.

**WSD spells out its own order: warmup, stable, decay.** Each phase has a job, and the data swap is timed to the transition into decay, not the start of the run. - **Phase 1 - warmup.** LR climbs linearly from near zero over the first few thousand steps, often paired with a batch-size ramp. The point is to keep update magnitudes inside the trust region while the weights settle out of the chaotic init region. - **Phase 2 - stable.** LR sits at its constant peak value over the bulk of the token budget, sometimes 80 to 95 percent of the run. The data mixture is the standard broad mix. The constant LR makes mid-run checkpoints comparable: any of them can be used as a starting point for a short decay branch. - **Phase 3 - anneal trigger.** As the decay window opens, the data mixture shifts toward scarce high-quality sources: curated textbooks, vetted code, exam-style content, instruction data. The shift coincides with the start of decay because high-quality data has the most impact when the LR is dropping and the model is consolidating. - **Phase 4 - decay.** LR drops rapidly toward zero (often linear or cosine over 5 to 20 percent of total steps). The model integrates the premium mix into the trunk's representations. - **Phase 5 - final checkpoint.** The fully decayed model is evaluated and handed to post-training (SFT, preference tuning). The most common ordering error is placing the anneal at the start of the stable phase. That wastes the premium data on a model that is not yet in consolidation mode.

TL;DR

Warmup ramps the LR up, the stable phase holds peak LR over most of the budget, an anneal trigger swaps in premium data, the decay phase ramps LR down, then the final checkpoint goes to post-training.

Memory aid

Easy to grasp

Think of training a long-distance runner over several months. You start with easy jogs that get harder week by week so the runner adapts; that is warmup. Then comes the long block of full-intensity training at a steady pace; that is the stable phase. As the big race approaches you switch from generic miles to race-specific drills and quality workouts; that is the anneal data swap. In the final weeks you taper down the intensity to let the body absorb the work; that is the decay phase. Race day is the final checkpoint where you hand the runner off to the event. The order is fixed because each phase prepares the body for the next, and rearranging them would either burn the runner out or leave them undertrained.

Key concepts

Pretraining

Concept explanation~2 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

WSD, short for warmup-stable-decay, is the LR schedule that displaced cosine as the default at frontier scale. The recipe is named after its three LR phases, but in practice a complete WSD run has five visible phases: warmup, stable, anneal trigger, decay, and final checkpoint handoff. The ordering is not arbitrary; each transition is timed to a specific change in the optimization regime, and reordering breaks the recipe in predictable ways.

The question asks for the order. The answer is warmup, then stable, then anneal trigger, then decay, then final checkpoint. The two non-obvious points are that batch-size ramps live inside warmup rather than at the start of stable, and that the premium data swap is timed to the decay transition rather than to the start of the run or the start of stable.

This deep dive walks through each phase, what it does, why it is where it is, and what goes wrong if it is moved.

Warmup: get out of the chaotic init region safely

The first phase ramps LR linearly from near zero to peak over the first few thousand steps. Batch size is often ramped during this same window: the run starts at a smaller global batch and grows to the target batch as the LR rises.

Both ramps exist for the same reason. At initialization, the loss surface is high-curvature and Adam's moment estimates have not converged. Large updates at this stage push the weights outside the optimizer's trust region and the run diverges. Holding LR small until the weights move into a better-conditioned basin keeps the trajectory stable.

The warmup phase is short, typically 2000 to 8000 steps depending on model scale, but it is load-bearing. Compressing it (the bug in the warmup-bug question) reliably produces gradient explosion within the first few hundred steps.

Stable: the long plateau that makes branching possible

Anneal trigger and decay: where premium data lands

Final checkpoint: handoff to post-training

Situations where this technique stops working.

2–4 min · Everything important, quickly.

Real products, models, and research that use this idea.

MiniCPM from MBZUAI publishes a WSD recipe with explicit anneal-phase data-mixture changes
OLMo 2 from AI2 uses a WSD-style schedule with a documented anneal data swap into decay

What an interviewer would ask next. Try answering before peeking at the approach.

QWhy is the anneal data swap timed to the decay transition specifically?

Discuss consolidation regime, scarcity of premium data, and the empirical finding that quality-data impact compounds when LR is dropping rather than at peak LR.

1 more follow-up an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

Placing the anneal data swap at the start of the stable phase rather than at the transition into decay, missing the point that the premium mix is consumed while LR is dropping.

60 second bullets to scan on the way to the call.

What each WSD phase is responsible for
Why batch-size ramps live in warmup rather than stable

Primary sources. Browse if you want the original framing.