Learning Rate Scheduling - PyTorch Basics

Learning Rate Scheduling

What you'll build: lr_schedule(step, warmup_steps, total_steps, base_lr, min_lr), the learning rate at a given step under linear warmup followed by cosine decay.

Background

A fixed learning rate is rarely best. Early in training, Adam's second-moment estimates are noisy and large updates can destabilise the model, so the rate is warmed up from 0. Later, smaller steps let the model settle into a minimum, so the rate is decayed. Warmup + cosine decay is the standard schedule for transformers (GPT-3, LLaMA, ViT).

Formula

if step < warmup_steps:                       # linear warmup, 0 -> base_lr
    lr = base_lr * step / warmup_steps
else:                                         # cosine decay, base_lr -> min_lr
    progress = (step - warmup_steps) / (total_steps - warmup_steps)    # 0 -> 1
    lr = min_lr + 0.5 * (base_lr - min_lr) * (1 + cos(π * progress))

At progress = 0 the cosine term is 1 (lr = base_lr), at 0.5 it is 0 (half way between the two), at 1 it is -1 (lr = min_lr).

Inputs

• step: integer in [0, total_steps].
• warmup_steps: integer in [0, total_steps); 0 means no warmup.
• total_steps: integer, the step at which the schedule reaches min_lr.
• base_lr: peak rate (default 1e-3); min_lr: final rate (default 1e-6).

Expected Return Format

A Python float.

Edge Cases and Conventions

• step == warmup_steps belongs to the decay branch and returns exactly base_lr, so the two pieces join continuously.
• step == 0 returns 0.0 when there is warmup.

Constraints

• PyTorch track: return a plain float (use math.cos). The function plugs into torch.optim.lr_scheduler.LambdaLR, which the hidden tests use.

Follow-up

• With warmup_steps = 0 this equals PyTorch's CosineAnnealingLR. Check it.
• What does "warmup-stable-decay" (WSD) change, and why do some LLM runs prefer it?

Input: [lr_schedule(s, warmup_steps=100, total_steps=1000) for s in (0, 50, 100)]
Output: [0.0, 0.0005, 0.001]

Input: [lr_schedule(s, 100, 1000) for s in (550, 1000)]
Output: [0.0005005, 1e-06]

import numpy as np

def lr_schedule(step, warmup_steps, total_steps, base_lr=0.001, min_lr=1e-6):
    """
    Learning rate with linear warmup and cosine decay.

    Args:
        step: Current training step (0 <= step <= total_steps)
        warmup_steps: Number of warmup steps (0 for none)
        total_steps: Step at which the rate reaches min_lr
        base_lr: Peak learning rate (default 0.001)
        min_lr: Final learning rate (default 1e-6)

    Returns:
        Learning rate (float) for the given step
    """
    # Your code here
    pass