| """WSD (warmup-stable-decay / trapezoidal) LR schedule (plan §6). |
| |
| Flat peak in the middle so we can branch/anneal at any point (seed-and-soup, epoch probe). |
| Decay phase = the gold-only anneal window. Returns a multiplier in [0,1] of peak LR. |
| """ |
| from __future__ import annotations |
|
|
| import math |
|
|
|
|
| def wsd(step, total, warmup_frac=0.04, decay_frac=0.2, min_frac=0.02): |
| warm = int(total * warmup_frac) |
| decay_start = int(total * (1 - decay_frac)) |
| if step < warm: |
| return step / max(warm, 1) |
| if step < decay_start: |
| return 1.0 |
| |
| t = (step - decay_start) / max(total - decay_start, 1) |
| return min_frac + (1 - min_frac) * 0.5 * (1 + math.cos(math.pi * t)) |
|
|
|
|
| def in_anneal(step, total, decay_frac=0.2): |
| """True once we enter the decay window (loader switches to gold-only).""" |
| return step >= int(total * (1 - decay_frac)) |
|
|
|
|
| def wsd_dyn(step, total, anneal_start, warmup_frac=0.04, decay_frac=0.2, min_frac=0.02): |
| """WSD with a dynamic anneal point. anneal_start <= the planned decay start; the decay |
| window keeps its planned LENGTH (decay_frac * total), so an early anneal finishes the |
| run early rather than stretching the decay. With anneal_start == planned start this is |
| identical to wsd().""" |
| warm = int(total * warmup_frac) |
| decay_len = total - int(total * (1 - decay_frac)) |
| if step < warm: |
| return step / max(warm, 1) |
| if step < anneal_start: |
| return 1.0 |
| t = min((step - anneal_start) / max(decay_len, 1), 1.0) |
| return min_frac + (1 - min_frac) * 0.5 * (1 + math.cos(math.pi * t)) |
|
|
