Warmup + Cosine Learning-Rate Schedule

The problem

LLMs (large language models) are not trained with a fixed learning rate. It changes at every step, following a schedule: it grows in a straight line during a short warmup, then follows half a cosine wave down to a minimum, and stays there. The theory page explains why each phase helps.

You will write the function that a training loop calls once per step to get the learning rate: lr_at(step, max_lr, min_lr, warmup_steps, total_steps), with steps counted from 0.

In the formulas below:

• $s$ is step, $W$ is warmup_steps, $S$ is total_steps;
• $\eta_{\max}$ (eta max) is max_lr and $\eta_{\min}$ is min_lr.

You can assume total_steps > warmup_steps >= 0.

Phase 1: warmup, when step < warmup_steps

The rate rises in $W$ equal increments, so that step 0 gets $\eta_{\max} / W$ and step $W - 1$ gets exactly $\eta_{\max}$:

Example. $W = 4$ and $\eta_{\max} = 0.001$: steps 0, 1, 2, 3 get $0.00025$, $0.0005$, $0.00075$ and $0.001$.

Phase 2: cosine decay, when warmup_steps <= step <= total_steps

where:

• $r = \frac{s - W}{S - W}$ is the progress through the decay: 0 on the first step after warmup, 1 on the last step;
• $\cos \pi r$ goes from $1$ (at $r = 0$) to $-1$ (at $r = 1$), so the bracket $1 + \cos \pi r$ goes from 2 to 0;
• so $\eta$ slides from $\eta_{\max}$ down to $\eta_{\min}$. Halfway, at $r = 0.5$, $\cos \tfrac{\pi}{2} = 0$ and $\eta$ is exactly the average of the two.

Phase 3: after the end, when step > total_steps

Stay at $\eta_{\min}$.

Example. The whole schedule with $W = 2$, $S = 10$, $\eta_{\max} = 0.001$, $\eta_{\min} = 0.0001$:

step	phase	$r$	learning rate
0	warmup		0.0005
1	warmup		0.001
2	cosine	0	0.001
4	cosine	0.25	0.000868
6	cosine	0.5	0.00055
8	cosine	0.75	0.000232
10	cosine	1	0.0001
11	after		0.0001

Step 6 is halfway through the decay: $(0.001 + 0.0001) / 2 = 0.00055$.

What to implement

lr_at(step, max_lr, min_lr, warmup_steps, total_steps): the learning rate of step step, as a Python float.

Requirements

• warmup_steps = 0 means no warmup: step 0 is already in phase 2, with $r = 0$, and gets max_lr. Do not divide by zero.
• Use math.cos and math.pi, no loops: this function is called once per training step.