| """Learning-rate and checkpoint schedules (token-indexed). |
| |
| WSD (warmup-stable-decay): a long constant-LR trunk lets us checkpoint a |
| `stable` point and later branch a cheap cooldown to any token budget without |
| retraining the trunk -- the mechanism the plan uses to extend 4 runs to 100B. |
| |
| All schedules are pure functions of tokens-seen-in-this-trajectory, so a |
| cooldown branch is just a run whose schedule has warmup=stable=0. |
| """ |
| import math |
|
|
|
|
| def lr_at(t: float, sched: dict) -> float: |
| peak, mn = sched["peak_lr"], sched.get("min_lr", 0.0) |
| w = sched.get("warmup_tokens", 0) |
| s = sched.get("stable_tokens", 0) |
| d = sched.get("decay_tokens", 0) |
| if t < w: |
| return peak * t / max(w, 1) |
| if t < w + s: |
| return peak |
| if t < w + s + d: |
| prog = (t - w - s) / max(d, 1) |
| shape = sched.get("decay_shape", "1-sqrt") |
| if shape == "linear": |
| f = 1.0 - prog |
| elif shape == "cosine": |
| f = 0.5 * (1.0 + math.cos(math.pi * prog)) |
| else: |
| f = 1.0 - math.sqrt(prog) |
| return mn + (peak - mn) * f |
| return mn |
|
|
|
|
| def total_tokens(sched: dict) -> float: |
| return (sched.get("warmup_tokens", 0) + sched.get("stable_tokens", 0) |
| + sched.get("decay_tokens", 0)) |
|
|
|
|
| def stable_end_tokens(sched: dict) -> float: |
| """Token count at which the trunk's decay begins (branch point).""" |
| return sched.get("warmup_tokens", 0) + sched.get("stable_tokens", 0) |
|
|
|
|
| def ckpt_interval(tokens: float, table: list) -> float: |
| """Log-spaced checkpoint interval: dense early, sparse late.""" |
| for up_to, interval in table: |
| if tokens < up_to: |
| return interval |
| return table[-1][1] |
|
|
