49. Warmup + Cosine Learning Rate Scheduler

Problem

Training large models from a random or freshly-loaded checkpoint with the full target learning rate from step 0 is unstable -- gradients are large and noisy early on. The standard fix is a linear warmup (ramp the LR up from 0 to the target base_lr over the first warmup_steps steps) followed by a cosine decay down to a floor min_lr over the rest of training. This shape -- linear ramp up, smooth cosine ramp down -- is close to the most common LR schedule used to pretrain transformer-scale models.

Implement a vectorized LR-schedule function: given a tensor of step indices, return the learning rate at each of those steps.

Source: src/49_warmup_cosine_lr_scheduler.py

def warmup_cosine_lr(
    steps: torch.Tensor, base_lr: float, warmup_steps: int,
    total_steps: int, min_lr: float = 0.0,
) -> torch.Tensor: ...

>>> steps = torch.tensor([0, 5, 10, 20])
>>> warmup_cosine_lr(steps, base_lr=1.0, warmup_steps=10, total_steps=20)
tensor([0.0000, 0.5000, 1.0000, 0.0000])

Step-by-Step Approach

1. Cast steps to float so the arithmetic below produces a float tensor rather than truncating on integer division.
2. Compute the linear-warmup value for every step, unconditionally: base_lr * step / warmup_steps. It's fine that this formula is "wrong" for steps past warmup -- it gets discarded by the selection step later.
3. Compute progress, the fraction of the way through the post-warmup decay: (step - warmup_steps) / (total_steps - warmup_steps), clamped to [0, 1] so steps beyond total_steps stay pinned at the end of the decay instead of the cosine curve rising back up.
4. Compute the cosine-decay value: min_lr + (base_lr - min_lr) * 0.5 * (1 + cos(pi * progress)) -- at progress=0 this equals base_lr; at progress=1 it equals min_lr.
5. Select per-element between the two branches with torch.where(step < warmup_steps, warmup_lr, cosine_lr).
6. Verify the boundary: at step == warmup_steps, progress == 0 so cosine_lr == base_lr, matching what the warmup branch would have given at that same step -- the two pieces meet with no discontinuity.

The key insight is computing both branches unconditionally and vectorized with torch.where, rather than branching in Python per step -- this makes the function work directly on a whole tensor of step indices at once, and avoids any risk of a Python-level if desyncing from the tensor's actual dtype or device.

Reference solution

def warmup_cosine_lr(
    steps: torch.Tensor,
    base_lr: float,
    warmup_steps: int,
    total_steps: int,
    min_lr: float = 0.0,
) -> torch.Tensor:
    steps_f = steps.to(torch.float32)

    # Linear warmup branch: computed for EVERY step (cheap, vectorized),
    # even the ones that will end up using the cosine branch instead --
    # torch.where below picks per-element which value survives.
    warmup_lr = base_lr * steps_f / warmup_steps

    # clamp to [0, 1] so steps past total_steps hold flat at min_lr instead
    # of the cosine continuing past its trough and rising again.
    progress = torch.clamp(
        (steps_f - warmup_steps) / (total_steps - warmup_steps), 0.0, 1.0
    )
    cosine_lr = min_lr + (base_lr - min_lr) * 0.5 * (1.0 + torch.cos(math.pi * progress))

    # steps_f < warmup_steps selects warmup_lr; the cosine branch handles
    # everything else, including step == warmup_steps exactly (no
    # discontinuity at the handoff).
    return torch.where(steps_f < warmup_steps, warmup_lr, cosine_lr)

Key Functions & Tricks

• torch.where(cond, a, b) — elementwise selection between two precomputed tensors, the standard way to vectorize a per-element branch instead of looping in Python.
• torch.clamp(t, 0.0, 1.0) — pins progress into a valid range so the cosine formula never runs past its intended domain.
• torch.cos — elementwise cosine; combined with math.pi, implements the smooth decay curve.
• steps.to(torch.float32) — avoids integer-division truncation when steps arrives as an int64 tensor.

How to Recognize This Pattern