Learning Rate Schedules: Warmup, Cosine Decay, and Why WSD Changes Everything

The learning rate is the single most important hyperparameter in LLM training. Too high: training diverges (loss spikes to infinity). Too low: training converges but wastes compute by moving too slowly. The schedule — how the learning rate changes over training — determines stability, final quality, and compute efficiency.

Why Not a Constant Learning Rate?

A constant LR fails for transformers because the loss landscape changes during training:

• Early training: The model is far from any minimum. Large gradients dominate. A high LR causes overshooting and instability, especially with randomly initialized attention weights that produce near-uniform attention distributions (high-entropy softmax outputs create large gradients through the attention mechanism).
• Mid training: The model is in a reasonable basin. Moderate LR allows efficient descent.
• Late training: The model is near a minimum. High LR causes oscillation around the minimum rather than convergence. Low LR is needed for fine-grained optimization.

Linear Warmup

The first 0.1-2% of training uses a linearly increasing LR:

$\text{lr}(t) = \text{lr}_{\max} \times \frac{t}{T_{\text{warmup}}}$

def linear_warmup(step, max_lr, warmup_steps):
    if step < warmup_steps:
        return max_lr * step / warmup_steps
    return max_lr

Why warmup exists: At step 0, all weights are randomly initialized. The Adam optimizer’s second moment estimates ($v_t$) are zero — they haven’t accumulated any gradient history. Without warmup, the first few updates use $\text{lr} / \sqrt{v_t + \epsilon}$ with tiny $v_t$, producing extremely large effective updates that can destabilize training.

Warmup gives Adam time to estimate the gradient statistics before applying full-strength updates. Typical warmup: 2,000 steps (Llama 3), which is 0.1% of the total 2M training steps.

Cosine Decay

After warmup, the standard schedule decays the LR following a cosine curve to a minimum value:

$\text{lr}(t) = \text{lr}_{\min} + \frac{1}{2}(\text{lr}_{\max} - \text{lr}_{\min})\left(1 + \cos\left(\frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}} \pi\right)\right)$

import math

def cosine_schedule(step, max_lr, min_lr, warmup_steps, total_steps):
    if step < warmup_steps:
        return max_lr * step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))

Why cosine? The cosine curve has two useful properties: (1) it starts with a gentle decline (near the maximum for the first ~25% of decay), allowing most training to happen at high LR, and (2) it has a long tail at low LR, allowing fine-grained convergence. Empirically, cosine outperforms linear decay by 0.5-1% on downstream benchmarks.

Typical minimum LR: $\text{lr}_{\min} = 0.1 \times \text{lr}_{\max}$ (Llama 3 uses this ratio).

WSD: Warmup-Stable-Decay

A newer schedule gaining adoption (used in some DeepSeek and Qwen training runs):

1. Warmup: Linear increase to $\text{lr}_{\max}$ (same as before)
2. Stable: Hold at $\text{lr}_{\max}$ for most of training (60-80% of steps)
3. Decay: Rapid decay (linear or cosine) in the final 20-40% of steps

def wsd_schedule(step, max_lr, min_lr, warmup_steps, stable_steps, total_steps):
    if step < warmup_steps:
        return max_lr * step / warmup_steps
    elif step < warmup_steps + stable_steps:
        return max_lr
    else:
        decay_steps = total_steps - warmup_steps - stable_steps
        decay_progress = (step - warmup_steps - stable_steps) / decay_steps
        return min_lr + (max_lr - min_lr) * (1 - decay_progress)

Why WSD? Cosine decay starts reducing the LR immediately after warmup. This means most training happens at a reduced LR. WSD keeps the LR at maximum for longer, allowing the model to explore more of the loss landscape before settling. The rapid final decay then focuses the model into a sharp minimum.

Peak Learning Rate Selection

The maximum LR depends on model size. Larger models need lower LR (their gradients have more variance):

The rough scaling: $\text{lr}_{\max} \propto N^{-0.5}$ where $N$ is parameter count. This is an empirical observation, not a derived formula. mu-P (Transformer Anatomy Part 17) provides a principled alternative.

Complete Implementation

class LRScheduler:
    """Configurable LR scheduler supporting warmup + cosine or WSD."""

    def __init__(self, max_lr, min_lr, warmup_steps, total_steps,
                 schedule="cosine", stable_fraction=0.0):
        self.max_lr = max_lr
        self.min_lr = min_lr
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.schedule = schedule
        self.stable_steps = int(stable_fraction * total_steps)

    def get_lr(self, step):
        # Warmup phase
        if step < self.warmup_steps:
            return self.max_lr * step / self.warmup_steps

        if self.schedule == "wsd":
            # Stable phase
            if step < self.warmup_steps + self.stable_steps:
                return self.max_lr
            # Decay phase
            decay_start = self.warmup_steps + self.stable_steps
            decay_total = self.total_steps - decay_start
            progress = (step - decay_start) / max(1, decay_total)
            return self.min_lr + (self.max_lr - self.min_lr) * (1 - progress)

        else:  # cosine
            progress = (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
            import math
            return self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (
                1 + math.cos(math.pi * progress)
            )

# Usage for Llama 3 70B:
scheduler = LRScheduler(
    max_lr=1.5e-4, min_lr=1.5e-5,
    warmup_steps=2000, total_steps=2_000_000,
    schedule="cosine"
)

Reviewer Agent Validation

Challenge: Compute the learning rate at step 500,000 for a Llama 3 70B training run with cosine schedule (max_lr=1.5e-4, min_lr=1.5e-5, warmup=2000, total=2M steps).

Expected: Progress = (500000 - 2000) / (2000000 - 2000) = 0.2494. LR = 1.5e-5 + 0.5 * (1.5e-4 - 1.5e-5) * (1 + cos(pi * 0.2494)) = 1.5e-5 + 6.75e-5 * (1 + cos(0.784)) = 1.5e-5 + 6.75e-5 * 1.707 = 1.5e-5 + 1.152e-4 = 1.302e-4.