| """Stage 2 settings, learning-rate schedule and checkpoint protection. |
| |
| Pure Python (no torch), so every rule here can be unit-tested on a laptop. |
| Stage 2 is switched on only by environment variables (exported from run_args.sh); |
| when STAGE2_START is not set, everything behaves exactly like Stage 1. |
| """ |
| import math |
| import os |
| from dataclasses import dataclass |
|
|
| |
| EOT_ID = 151645 |
|
|
|
|
| @dataclass |
| class Stage2Settings: |
| start: int |
| warmup: int |
| peak: float |
| min_lr: float |
| bins: list |
| weights: list |
|
|
| @staticmethod |
| def from_env(): |
| """Read Stage 2 settings from the environment; None means 'no Stage 2'.""" |
| if not os.environ.get("STAGE2_START"): |
| return None |
| bins = os.environ["S2_BIN"].split() |
| weights = [float(w) for w in os.environ["S2_WEIGHTS"].split()] |
| assert len(bins) == len(weights) > 0, "S2_BIN and S2_WEIGHTS must have the same length" |
| assert all(w >= 0 for w in weights) and sum(weights) > 0, "bad S2_WEIGHTS" |
| s = Stage2Settings( |
| start=int(os.environ["STAGE2_START"]), |
| warmup=int(os.environ.get("STAGE2_WARMUP", "1000")), |
| peak=float(os.environ.get("STAGE2_PEAK", "1e-4")), |
| min_lr=float(os.environ.get("STAGE2_MIN", "1e-5")), |
| bins=bins, |
| weights=weights, |
| ) |
| assert s.start > 0 and s.warmup > 0 and s.peak > s.min_lr > 0, "bad Stage 2 settings" |
| return s |
|
|
|
|
| def stage1_lr(step, cfg): |
| """Stage 1 schedule, unchanged: linear warmup, then cosine to min_lr at lr_horizon.""" |
| if step < cfg.warmup_steps: |
| return cfg.peak_lr * (step + 1) / cfg.warmup_steps |
| horizon = getattr(cfg, "lr_horizon", cfg.max_steps) |
| progress = min(1.0, (step - cfg.warmup_steps) / max(1, horizon - cfg.warmup_steps)) |
| return cfg.min_lr + 0.5 * (1 + math.cos(math.pi * progress)) * (cfg.peak_lr - cfg.min_lr) |
|
|
|
|
| def lr_at(step, cfg, s2=None): |
| """Stage 1 schedule before s2.start; after it: warm up from Stage 1's final LR |
| to s2.peak, then cosine down to s2.min_lr at the last step (cfg.max_steps).""" |
| if s2 is None or step < s2.start: |
| return stage1_lr(step, cfg) |
| t = step - s2.start |
| base = stage1_lr(s2.start - 1, cfg) |
| if t < s2.warmup: |
| return base + (s2.peak - base) * (t + 1) / s2.warmup |
| span = max(1, cfg.max_steps - 1 - s2.start - s2.warmup) |
| progress = min(1.0, (t - s2.warmup) / span) |
| return s2.min_lr + 0.5 * (1 + math.cos(math.pi * progress)) * (s2.peak - s2.min_lr) |
|
|
|
|
| def protected_ckpts(): |
| """Checkpoint folder names that rotation must never delete (env PROTECT_CKPTS).""" |
| return set(os.environ.get("PROTECT_CKPTS", "").split()) |
|
|
|
|
| def ckpts_to_delete(entries, is_complete, keep, protect): |
| """Which ckpt_<step> folders rotation removes: all half-written ones, plus complete |
| ones older than the newest `keep`. Protected names are never removed or counted.""" |
| entries = [d for d in entries if d.startswith("ckpt_") and d not in protect] |
| good = sorted((d for d in entries if is_complete(d)), key=lambda d: int(d.split("_")[1])) |
| bad = [d for d in entries if not is_complete(d)] |
| return (good[:-keep] if keep > 0 else good) + bad |
|
|
