Learning Rate: Warmup + Cosine Decay

Almost every modern LLM is trained with a linear warmup followed by cosine decay: start small so early updates don't blow up, peak, then glide down.

Implement lr_at(step, max_lr, warmup_steps, total_steps, min_lr=0.0):

• Warmup (step < warmup_steps): rise linearly, max_lr * (step + 1) / warmup_steps.
• Decay (warmup_steps <= step < total_steps): with progress $p = \frac{\text{step} - \text{warmup}}{\text{total} - \text{warmup}}$,

• After total_steps: stay at min_lr.