File size: 1,627 Bytes
5952424 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """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
# cosine decay to min_frac over the decay window
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))
|