Upload src/xscript/schedule.py with huggingface_hub

src/xscript/schedule.py

@@ -0,0 +1,50 @@
+"""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:  # "1-sqrt" (MiniCPM WSD; empirically strong)
+            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]