| """Learning-rate schedule: linear warmup, then cosine decay to a floor.""" | |
| from __future__ import annotations | |
| import math | |
| def lr_at(step: int, peak_lr: float, warmup_steps: int, total_steps: int, min_ratio: float = 0.1) -> float: | |
| if step < warmup_steps: | |
| return peak_lr * (step + 1) / warmup_steps | |
| progress = min(1.0, (step - warmup_steps) / max(1, total_steps - 1 - warmup_steps)) | |
| floor = peak_lr * min_ratio | |
| return floor + (peak_lr - floor) * 0.5 * (1 + math.cos(math.pi * progress)) | |