Spaces:

juiceb0xc0de
/

lr-scheduler-benchmark

Sleeping

lr-scheduler-benchmark / benchmark_runner.py
add 6 new schedulers: step_decay, cosine_annealing, plateau, cyclic, mnist, sqrt
aab71c5 verified 7 months ago
36.5 kB
"""
LR Scheduler Benchmark Runner
==============================
Benchmarks LR schedulers under identical conditions:
  - Model: distilbert-base-uncased (67M params)
  - Task: SST-2 sentiment classification (GLUE)
  - Optimizer: AdamW (lr=2e-5, weight_decay=0.01, betas=(0.9, 0.999))
  - Batch size: 32
  - Epochs: 3
  - Warmup: 6% of total steps
  - Seeds: 3 (report mean ± std)
Only the LR scheduler changes between runs. Everything else is frozen.
Outputs:
  results/leaderboard.csv      — one row per scheduler (aggregated across seeds)
  results/curves.csv           — per-step loss/lr curves for all runs
  results/per_run.csv          — per-seed results for each scheduler
Usage:
  python benchmark_runner.py                    # run all built-in schedulers
  python benchmark_runner.py --schedulers cosine linear wsd   # run specific ones
  python benchmark_runner.py --seeds 5          # more seeds for significance
"""

import argparse
import json
import math
import os
import time
from collections import deque
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
import torch
from torch.optim import AdamW
from torch.optim.lr_scheduler import (
    CosineAnnealingLR,
    OneCycleLR,
    ExponentialLR,
    StepLR,
    LambdaLR,
    CyclicLR,
    ReduceLROnPlateau,
    SequentialLR,
)
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    get_scheduler,
    DataCollatorWithPadding,
)
from torch.utils.data import DataLoader
from sklearn.metrics import accuracy_score


# ─── Configuration ────────────────────────────────────────────────────────────
MODEL_NAME = "distilbert/distilbert-base-uncased"
DATASET_NAME = "nyu-mll/glue"
DATASET_CONFIG = "sst2"
BATCH_SIZE = 32
NUM_EPOCHS = 3
LEARNING_RATE = 2e-5
WEIGHT_DECAY = 0.01
WARMUP_FRACTION = 0.06
EVAL_EVERY_STEPS = 50  # evaluate every N training steps
LOG_EVERY_STEPS = 10   # log train loss every N steps
TARGET_LOSS = 0.35     # threshold for "steps to target" metric
MAX_GRAD_NORM = 1.0
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ─── Scheduler Factories ─────────────────────────────────────────────────────
# Each factory takes (optimizer, num_warmup_steps, num_training_steps) and
# returns (scheduler, step_mode) where step_mode is:
#   "step"  — call scheduler.step() every step
#   "loss"  — call scheduler.step(loss=...) every step
#   "aecs"  — call scheduler.step(loss=..., grad_norm=...) every step

def make_linear(optimizer, warmup, total):
    s = get_scheduler("linear", optimizer, num_warmup_steps=warmup, num_training_steps=total)
    return s, "step"

def make_cosine(optimizer, warmup, total):
    s = get_scheduler("cosine", optimizer, num_warmup_steps=warmup, num_training_steps=total)
    return s, "step"

def make_cosine_with_restarts(optimizer, warmup, total):
    s = get_scheduler("cosine_with_restarts", optimizer, num_warmup_steps=warmup,
                      num_training_steps=total, scheduler_specific_kwargs={"num_cycles": 3})
    return s, "step"

def make_polynomial(optimizer, warmup, total):
    s = get_scheduler("polynomial", optimizer, num_warmup_steps=warmup, num_training_steps=total,
                      scheduler_specific_kwargs={"power": 2.0, "lr_end": 1e-7})
    return s, "step"

def make_constant(optimizer, warmup, total):
    s = get_scheduler("constant", optimizer, num_warmup_steps=0, num_training_steps=total)
    return s, "step"

def make_constant_with_warmup(optimizer, warmup, total):
    s = get_scheduler("constant_with_warmup", optimizer, num_warmup_steps=warmup,
                      num_training_steps=total)
    return s, "step"

def make_inverse_sqrt(optimizer, warmup, total):
    s = get_scheduler("inverse_sqrt", optimizer, num_warmup_steps=warmup,
                      num_training_steps=total)
    return s, "step"

def make_wsd(optimizer, warmup, total):
    num_decay = int(0.10 * total)
    num_stable = total - warmup - num_decay
    s = get_scheduler("wsd", optimizer, num_warmup_steps=warmup,
                      num_training_steps=total,
                      scheduler_specific_kwargs={
                          "num_stable_steps": num_stable,
                          "num_decay_steps": num_decay,
                      })
    return s, "step"

def make_onecycle(optimizer, warmup, total):
    s = OneCycleLR(optimizer, max_lr=LEARNING_RATE * 10, total_steps=total,
                   pct_start=0.3, anneal_strategy='cos')
    return s, "step"

def make_exponential(optimizer, warmup, total):
    # Decay so that final LR is ~10% of initial
    gamma = (0.1) ** (1.0 / max(1, total))
    # Add warmup via LambdaLR wrapper
    warmup_scheduler = LambdaLR(optimizer, lr_lambda=lambda step: min(1.0, step / max(1, warmup)))
    exp_scheduler = ExponentialLR(optimizer, gamma=gamma)
    from torch.optim.lr_scheduler import SequentialLR
    s = SequentialLR(optimizer, schedulers=[warmup_scheduler, exp_scheduler], milestones=[warmup])
    return s, "step"

def make_greedy_lr(optimizer, warmup, total):
    return GreedyLR(optimizer, base_lr=LEARNING_RATE, warmup_steps=warmup), "loss"

def make_dlrs(optimizer, warmup, total):
    return DLRSScheduler(optimizer, base_lr=LEARNING_RATE, warmup_steps=warmup), "loss"

def make_aecs(optimizer, warmup, total):
    cfg = AECSConfig(
        base_lr=LEARNING_RATE,
        warmup_steps=warmup,
        total_steps=total,
        verbose=False,
    )
    return EventControlScheduler(optimizer, cfg), "aecs"


def make_step_decay(optimizer, warmup, total):
    # lr × 0.1 at each third of training, with linear warmup
    # math: lr_t = base * 0.1^(floor(step / step_size))
    step_size = max(1, (total - warmup) // 3)
    warmup_sched = LambdaLR(optimizer, lr_lambda=lambda s: min(1.0, s / max(1, warmup)))
    step_sched = StepLR(optimizer, step_size=step_size, gamma=0.1)
    return SequentialLR(optimizer, schedulers=[warmup_sched, step_sched], milestones=[warmup]), "step"


def make_cosine_annealing(optimizer, warmup, total):
    # PyTorch CosineAnnealingLR: η_t = η_min + 0.5(η_max - η_min)(1 + cos(π·t/T))
    # floor at 1% of peak LR
    warmup_sched = LambdaLR(optimizer, lr_lambda=lambda s: min(1.0, s / max(1, warmup)))
    cos_sched = CosineAnnealingLR(optimizer, T_max=max(1, total - warmup), eta_min=LEARNING_RATE * 0.01)
    return SequentialLR(optimizer, schedulers=[warmup_sched, cos_sched], milestones=[warmup]), "step"


def make_plateau_reduction(optimizer, warmup, total):
    # ReduceLROnPlateau: lr × 0.5 after `patience` steps with no improvement
    return _PlateauWrapper(optimizer, base_lr=LEARNING_RATE, warmup_steps=warmup), "loss"


def make_cyclic_lr(optimizer, warmup, total):
    # CLR triangular2: cycles between base_lr*0.1 and base_lr, amplitude halves each cycle
    # math: lr oscillates, peak decays as 1/2^cycle
    step_size = max(1, (total - warmup) // 10)  # ~5 full cycles
    return CyclicLR(optimizer,
                    base_lr=LEARNING_RATE * 0.1,
                    max_lr=LEARNING_RATE,
                    step_size_up=step_size,
                    mode="triangular2",
                    cycle_momentum=False), "step"


def make_mnist_example(optimizer, warmup, total):
    # Classic PyTorch MNIST recipe: StepLR gamma=0.7 per epoch
    # math: lr_t = base * 0.7^epoch
    steps_per_epoch = max(1, total // 3)  # benchmark runs 3 epochs
    return StepLR(optimizer, step_size=steps_per_epoch, gamma=0.7), "step"


def make_sqrt_scheduler(optimizer, warmup, total):
    # Square-root decay from "Attention Is All You Need"
    # math: lr = base * min(1, warmup/step)^0.5  (peaks at warmup, then ∝ 1/√step)
    def _sqrt_lambda(step):
        if step < warmup:
            return step / max(1, warmup)
        return math.sqrt(warmup / max(warmup, step))
    return LambdaLR(optimizer, lr_lambda=_sqrt_lambda), "step"


class _PlateauWrapper:
    """ReduceLROnPlateau adapter to fit the step_mode='loss' interface."""
    def __init__(self, optimizer, base_lr, warmup_steps, factor=0.5, patience=10):
        self.optimizer = optimizer
        self.warmup_steps = max(1, warmup_steps)
        self.base_lr = base_lr
        self._step = 0
        self.inner = ReduceLROnPlateau(optimizer, mode="min", factor=factor,
                                       patience=patience, min_lr=base_lr * 0.01)

    def step(self, loss=None):
        self._step += 1
        if self._step <= self.warmup_steps:
            lr = self.base_lr * self._step / self.warmup_steps
            for pg in self.optimizer.param_groups:
                pg["lr"] = lr
        elif loss is not None:
            self.inner.step(loss)

    def get_last_lr(self):
        return [pg["lr"] for pg in self.optimizer.param_groups]


def make_rex(optimizer, warmup, total):
    # REX: Reflected Exponential — stays high longer, drops sharply near end
    # Budhraja et al. 2023, "REX: Revisiting Budgeted Training with an Improved Schedule"
    alpha = 5.0
    def lr_lambda(step):
        if step < warmup:
            return step / max(1, warmup)
        r = min((step - warmup) / max(1, total - warmup), 1.0)
        return (math.exp(alpha * (1 - r)) - 1) / (math.exp(alpha) - 1)
    return LambdaLR(optimizer, lr_lambda=lr_lambda), "step"


def make_exp_increasing(optimizer, warmup, total):
    # Exponentially Increasing LR — LR grows exponentially through training
    # Useful for identifying optimal LR region; passes through sweet spot
    # lr sweeps from base_lr * 0.01 → base_lr * 5 over full run
    lr_min_frac = 0.01
    lr_max_frac = 5.0
    log_ratio = math.log(lr_max_frac / lr_min_frac)
    def lr_lambda(step):
        t = step / max(1, total)
        return lr_min_frac * math.exp(log_ratio * t)
    return LambdaLR(optimizer, lr_lambda=lr_lambda), "step"


def make_scaling_law_opt(optimizer, warmup, total):
    # Power-law decay derived from neural scaling law analysis
    # Optimal schedule under fixed compute budget: lr ∝ (1 - t/T)^β
    # β=0.9 from Kaplan et al. scaling law analysis; floor at 1% of peak
    beta = 0.9
    min_frac = 0.01
    def lr_lambda(step):
        if step < warmup:
            return step / max(1, warmup)
        r = min((step - warmup) / max(1, total - warmup), 1.0 - 1e-6)
        return max(min_frac, (1 - r) ** beta)
    return LambdaLR(optimizer, lr_lambda=lr_lambda), "step"


SCHEDULER_REGISTRY = {
    "linear":               make_linear,
    "cosine":               make_cosine,
    "cosine_with_restarts": make_cosine_with_restarts,
    "polynomial":           make_polynomial,
    "constant":             make_constant,
    "constant_with_warmup": make_constant_with_warmup,
    "inverse_sqrt":         make_inverse_sqrt,
    "wsd":                  make_wsd,
    "onecycle":             make_onecycle,
    "exponential":          make_exponential,
    "greedy_lr":            make_greedy_lr,
    "dlrs":                 make_dlrs,
    "aecs":                 make_aecs,
    "rex":                  make_rex,
    "exp_increasing":       make_exp_increasing,
    "scaling_law_opt":      make_scaling_law_opt,
    "step_decay":           make_step_decay,
    "cosine_annealing":     make_cosine_annealing,
    "plateau_reduction":    make_plateau_reduction,
    "cyclic_lr":            make_cyclic_lr,
    "mnist_example":        make_mnist_example,
    "sqrt_scheduler":       make_sqrt_scheduler,
}


# ─── Loss-aware schedulers ────────────────────────────────────────────────────

class GreedyLR:
    def __init__(self, optimizer, base_lr, warmup_steps,
                 factor_up=1.5, factor_down=0.5, patience=5, threshold=1e-3,
                 min_lr=None, max_lr=None):
        self.optimizer = optimizer
        self.base_lr = base_lr
        self.warmup_steps = max(1, warmup_steps)
        self.factor_up = factor_up
        self.factor_down = factor_down
        self.patience = patience
        self.threshold = threshold
        self.min_lr = min_lr if min_lr is not None else base_lr / 100
        self.max_lr = max_lr if max_lr is not None else base_lr * 10
        self.step_count = 0
        self.best_loss = float("inf")
        self.bad_steps = 0
        self._set_lr(0.0)

    def _set_lr(self, lr):
        for g in self.optimizer.param_groups:
            g["lr"] = lr

    def get_last_lr(self):
        return [g["lr"] for g in self.optimizer.param_groups]

    def step(self, loss=None):
        self.step_count += 1
        if self.step_count <= self.warmup_steps:
            self._set_lr(self.base_lr * self.step_count / self.warmup_steps)
            return
        if loss is None:
            return
        cur = self.optimizer.param_groups[0]["lr"]
        improved = loss < self.best_loss * (1 - self.threshold)
        if improved:
            self.best_loss = loss
            self.bad_steps = 0
            new_lr = min(cur * self.factor_up, self.max_lr)
        else:
            self.bad_steps += 1
            if self.bad_steps >= self.patience:
                new_lr = max(cur * self.factor_down, self.min_lr)
                self.bad_steps = 0
            else:
                new_lr = cur
        self._set_lr(new_lr)


class DLRSScheduler:
    def __init__(self, optimizer, base_lr, warmup_steps,
                 window=20, factor_up=1.1, factor_down=0.7,
                 slope_fast=-1e-3, slope_slow=-1e-4,
                 min_lr=None, max_lr=None):
        self.optimizer = optimizer
        self.base_lr = base_lr
        self.warmup_steps = max(1, warmup_steps)
        self.window = window
        self.factor_up = factor_up
        self.factor_down = factor_down
        self.slope_fast = slope_fast
        self.slope_slow = slope_slow
        self.min_lr = min_lr if min_lr is not None else base_lr / 100
        self.max_lr = max_lr if max_lr is not None else base_lr * 10
        self.step_count = 0
        self.loss_buf = []
        self._set_lr(0.0)

    def _set_lr(self, lr):
        for g in self.optimizer.param_groups:
            g["lr"] = lr

    def get_last_lr(self):
        return [g["lr"] for g in self.optimizer.param_groups]

    def step(self, loss=None):
        self.step_count += 1
        if self.step_count <= self.warmup_steps:
            self._set_lr(self.base_lr * self.step_count / self.warmup_steps)
            return
        if loss is None:
            return
        self.loss_buf.append(float(loss))
        if len(self.loss_buf) > self.window:
            self.loss_buf.pop(0)
        cur = self.optimizer.param_groups[0]["lr"]
        if len(self.loss_buf) < self.window:
            return
        xs = np.arange(len(self.loss_buf), dtype=float)
        slope = float(np.polyfit(xs, self.loss_buf, 1)[0])
        if slope <= self.slope_fast:
            new_lr = cur
        elif slope <= self.slope_slow:
            new_lr = min(cur * self.factor_up, self.max_lr)
        else:
            new_lr = max(cur * self.factor_down, self.min_lr)
        self._set_lr(new_lr)


# ─── AECS: Adaptive Event-Control Scheduler ──────────────────────────────────

@dataclass
class AECSConfig:
    base_lr: float = 5e-5
    warmup_steps: int = 500
    total_steps: int = 10000
    loss_window: int = 50
    grad_window: int = 50
    redundancy_window: int = 20
    instability_z_thresh: float = 2.5
    loss_spike_ratio: float = 1.15
    redundancy_thresh: float = 0.98
    plateau_grad_norm_thresh: float = 1e-4
    recovery_lr_factor: float = 0.3
    recovery_momentum_factor: float = 0.5
    recovery_clip_boost: float = 1.5
    recovery_min_steps: int = 100
    recovery_max_steps: int = 1000
    explore_lr_factor: float = 1.5
    explore_noise_std: float = 1e-5
    event_persistence: int = 5
    cooldown_steps: int = 200
    reentry_grad_norm_tol: float = 0.1
    verbose: bool = True


class SignalBuffer:
    def __init__(self, window: int = 50):
        self.losses: deque = deque(maxlen=window)
        self.grad_norms: deque = deque(maxlen=window)
        self.grad_cosines: deque = deque(maxlen=window)
        self.layer_grad_norms: deque = deque(maxlen=window)
        self.steps: int = 0
        self._prev_grad_flat: Optional[torch.Tensor] = None

    def push(self, loss: float, grad_norm: float, layer_grad_norms=None):
        self.losses.append(loss)
        self.grad_norms.append(grad_norm)
        self.layer_grad_norms.append(layer_grad_norms or [])
        self.steps += 1

    def loss_min_recent(self, n: int = 10) -> float:
        if len(self.losses) == 0:
            return float('inf')
        return min(list(self.losses)[-n:])

    def grad_norm_ema(self, alpha: float = 0.97) -> Tuple[float, float]:
        if len(self.grad_norms) == 0:
            return 0.0, 1.0
        mu = self.grad_norms[0]
        var = 0.0
        for g in list(self.grad_norms)[1:]:
            mu = alpha * mu + (1 - alpha) * g
            var = alpha * var + (1 - alpha) * (g - mu) ** 2
        return mu, math.sqrt(max(var, 1e-8))

    def grad_norm_zscore(self) -> float:
        if len(self.grad_norms) < 5:
            return 0.0
        mu, sigma = self.grad_norm_ema()
        if sigma < 1e-8:
            return 0.0
        return (self.grad_norms[-1] - mu) / sigma

    def grad_norm_variance(self) -> float:
        if len(self.grad_norms) < 5:
            return 0.0
        vals = list(self.grad_norms)
        mean = sum(vals) / len(vals)
        return sum((x - mean) ** 2 for x in vals) / len(vals)

    def redundancy_score(self) -> float:
        if len(self.grad_cosines) < max(1, self.grad_cosines.maxlen // 2):
            return 0.0
        vals = list(self.grad_cosines)
        return sum(vals) / len(vals)


class EventControlScheduler:
    """Adaptive Event-Control Scheduler (AECS).
    Modes: BASELINE, RECOVERY, EXPLORE, STABILIZE
    Cosine backbone with event-driven mode multipliers.
    Call scheduler.step({"loss": ..., "grad_norm": ...}) each training step.
    """
    MODES = ["BASELINE", "RECOVERY", "EXPLORE", "STABILIZE"]

    def __init__(self, optimizer, config: Optional[AECSConfig] = None):
        self.optimizer = optimizer
        self.config = config or AECSConfig()
        self.buffer = SignalBuffer(window=max(self.config.loss_window, self.config.grad_window))
        self.mode: str = "BASELINE"
        self.mode_steps: int = 0
        self.total_steps: int = 0
        self.event_counter: Dict[str, int] = {m: 0 for m in self.MODES}
        self.transition_log: List[Dict] = []
        self.base_lrs = [g["lr"] for g in optimizer.param_groups]
        self.base_betas = [g.get("betas", (0.9, 0.999)) for g in optimizer.param_groups]
        self.base_weight_decays = [g.get("weight_decay", 0.0) for g in optimizer.param_groups]
        self._loss_ema: float = 0.0
        self._loss_ema_alpha: float = 0.95

    def step(self, signals: Dict[str, float] = None, loss=None, grad_norm=0.0) -> str:
        # Accept both dict form and keyword form for compatibility
        if signals is None:
            signals = {"loss": loss or 0.0, "grad_norm": grad_norm}

        _loss = signals.get("loss", 0.0)
        _grad_norm = signals.get("grad_norm", 0.0)

        if self.total_steps == 0:
            self._loss_ema = _loss
        else:
            self._loss_ema = self._loss_ema_alpha * self._loss_ema + (1 - self._loss_ema_alpha) * _loss

        self.buffer.push(_loss, _grad_norm)
        self.total_steps += 1
        self.mode_steps += 1

        event = self._detect_event(signals)
        if event:
            self._maybe_transition(event)

        lrs = self._compute_lrs()
        for group, lr in zip(self.optimizer.param_groups, lrs):
            group["lr"] = lr

        self._apply_mode_tweaks()
        return self.mode

    def _detect_event(self, signals) -> Optional[str]:
        cfg = self.config
        buf = self.buffer
        if buf.steps < cfg.event_persistence + 5:
            return None
        events = []
        z = buf.grad_norm_zscore()
        if z > cfg.instability_z_thresh:
            events.append("GRADIENT_SPIKE")
        recent_min = buf.loss_min_recent(n=10)
        if recent_min > 0 and buf.losses[-1] > recent_min * cfg.loss_spike_ratio:
            events.append("LOSS_SPIKE")
        redundancy = buf.redundancy_score()
        if redundancy > cfg.redundancy_thresh:
            events.append("REDUNDANT")
        grad_var = buf.grad_norm_variance()
        if grad_var > cfg.reentry_grad_norm_tol and z > 1.0:
            events.append("UNSTABLE")
        if len(buf.grad_norms) >= 10:
            recent_avg = sum(list(buf.grad_norms)[-10:]) / 10
            if recent_avg < cfg.plateau_grad_norm_thresh:
                events.append("PLATEAU")
        if not events:
            return None
        for p in ["GRADIENT_SPIKE", "LOSS_SPIKE", "UNSTABLE", "REDUNDANT", "PLATEAU"]:
            if p in events:
                return p
        return events[0]

    def _maybe_transition(self, event: str):
        cfg = self.config
        if self.mode_steps < cfg.cooldown_steps and self.mode != "BASELINE":
            return
        if self.mode == "RECOVERY":
            if self.mode_steps < cfg.recovery_min_steps:
                return
            if self.mode_steps >= cfg.recovery_max_steps:
                self._enter_mode("BASELINE", "recovery_max_duration")
                return
        target = self._event_to_mode(event)
        if target == self.mode:
            return
        if target == "BASELINE" and self.mode in ("RECOVERY", "STABILIZE"):
            if self.buffer.grad_norm_variance() > cfg.reentry_grad_norm_tol:
                return
        self._enter_mode(target, event)

    def _event_to_mode(self, event: str) -> str:
        return {
            "GRADIENT_SPIKE": "RECOVERY",
            "LOSS_SPIKE": "RECOVERY",
            "UNSTABLE": "STABILIZE",
            "REDUNDANT": "EXPLORE",
            "PLATEAU": "EXPLORE",
        }.get(event, "BASELINE")

    def _enter_mode(self, new_mode: str, cause: str):
        old_mode = self.mode
        self.mode = new_mode
        self.mode_steps = 0
        self.event_counter[new_mode] += 1
        self.transition_log.append({
            "step": self.total_steps, "from": old_mode, "to": new_mode,
            "cause": cause, "lr": self.optimizer.param_groups[0]["lr"],
        })
        if self.config.verbose:
            print(f"[AECS] Step {self.total_steps}: {old_mode} -> {new_mode} ({cause})")

    def _compute_lrs(self) -> List[float]:
        cfg = self.config
        step = self.total_steps
        if step < cfg.warmup_steps:
            backbone = step / max(cfg.warmup_steps, 1)
        else:
            progress = (step - cfg.warmup_steps) / max(cfg.total_steps - cfg.warmup_steps, 1)
            backbone = 0.5 * (1.0 + math.cos(math.pi * progress))
        mult = 1.0
        if self.mode == "RECOVERY":
            mult = cfg.recovery_lr_factor
        elif self.mode == "EXPLORE":
            mult = cfg.explore_lr_factor
        elif self.mode == "STABILIZE":
            mult = cfg.recovery_lr_factor * 0.8
        return [base * backbone * mult for base in self.base_lrs]

    def _apply_mode_tweaks(self):
        cfg = self.config
        for group, base_b, base_wd in zip(
            self.optimizer.param_groups, self.base_betas, self.base_weight_decays
        ):
            if "betas" in group:
                beta1, beta2 = base_b
                if self.mode == "RECOVERY":
                    group["betas"] = (beta1 * cfg.recovery_momentum_factor, beta2)
                else:
                    group["betas"] = (beta1, beta2)
            if "weight_decay" in group:
                if self.mode == "EXPLORE":
                    group["weight_decay"] = base_wd * 0.5
                elif self.mode == "RECOVERY":
                    group["weight_decay"] = base_wd * 1.2
                else:
                    group["weight_decay"] = base_wd

    def get_last_lr(self) -> List[float]:
        return [g["lr"] for g in self.optimizer.param_groups]

    def summary(self) -> Dict:
        return {
            "mode": self.mode,
            "total_steps": self.total_steps,
            "transitions": len(self.transition_log),
            "event_counter": self.event_counter,
            "recent_transitions": self.transition_log[-5:],
        }


# ─── Data Loading ─────────────────────────────────────────────────────────────

def load_data(tokenizer):
    dataset = load_dataset(DATASET_NAME, DATASET_CONFIG)

    def preprocess(examples):
        return tokenizer(examples["sentence"], truncation=True, max_length=128)

    tokenized = dataset.map(preprocess, batched=True, remove_columns=["sentence", "idx"])
    tokenized = tokenized.rename_column("label", "labels")
    tokenized.set_format("torch")

    collator = DataCollatorWithPadding(tokenizer=tokenizer)
    train_loader = DataLoader(tokenized["train"], batch_size=BATCH_SIZE, shuffle=True,
                              collate_fn=collator, num_workers=2, pin_memory=True)
    val_loader = DataLoader(tokenized["validation"], batch_size=BATCH_SIZE * 2,
                            collate_fn=collator, num_workers=2, pin_memory=True)
    return train_loader, val_loader


# ─── Evaluation ───────────────────────────────────────────────────────────────

@torch.no_grad()
def evaluate(model, val_loader):
    model.eval()
    total_loss = 0
    all_preds, all_labels = [], []
    for batch in val_loader:
        batch = {k: v.to(DEVICE) for k, v in batch.items()}
        outputs = model(**batch)
        total_loss += outputs.loss.item() * batch["labels"].size(0)
        preds = outputs.logits.argmax(dim=-1)
        all_preds.extend(preds.cpu().numpy())
        all_labels.extend(batch["labels"].cpu().numpy())
    avg_loss = total_loss / len(all_labels)
    acc = accuracy_score(all_labels, all_preds)
    model.train()
    return avg_loss, acc


# ─── Single Run ───────────────────────────────────────────────────────────────

def run_single(scheduler_name, seed, tokenizer, train_loader, val_loader):
    """Run one training session. Returns (metrics_dict, curve_records)."""
    torch.manual_seed(seed)
    np.random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_NAME, num_labels=2
    ).to(DEVICE)
    model.train()

    optimizer = AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY,
                      betas=(0.9, 0.999), eps=1e-8)

    num_training_steps = len(train_loader) * NUM_EPOCHS
    num_warmup_steps = int(WARMUP_FRACTION * num_training_steps)

    factory = SCHEDULER_REGISTRY[scheduler_name]
    scheduler, step_mode = factory(optimizer, num_warmup_steps, num_training_steps)

    val_loss, val_acc = evaluate(model, val_loader)

    curve_records = []
    step = 0
    running_loss = 0.0
    running_count = 0
    steps_to_target = None
    start_time = time.time()

    lr_now = optimizer.param_groups[0]["lr"]
    curve_records.append({
        "scheduler": scheduler_name, "seed": seed, "step": 0,
        "train_loss": None, "val_loss": val_loss, "val_acc": val_acc, "lr": lr_now
    })

    for epoch in range(NUM_EPOCHS):
        for batch in train_loader:
            batch = {k: v.to(DEVICE) for k, v in batch.items()}
            outputs = model(**batch)
            loss = outputs.loss
            loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
            optimizer.step()

            if step_mode == "aecs":
                scheduler.step({"loss": loss.item(), "grad_norm": grad_norm.item()})
            elif step_mode == "loss":
                scheduler.step(loss=loss.item())
            else:
                scheduler.step()

            optimizer.zero_grad()

            step += 1
            running_loss += loss.item()
            running_count += 1

            lr_now = optimizer.param_groups[0]["lr"]

            if step % LOG_EVERY_STEPS == 0:
                avg_train_loss = running_loss / running_count
                curve_records.append({
                    "scheduler": scheduler_name, "seed": seed, "step": step,
                    "train_loss": avg_train_loss, "val_loss": None, "val_acc": None,
                    "lr": lr_now
                })
                running_loss = 0.0
                running_count = 0

            if step % EVAL_EVERY_STEPS == 0:
                vl, va = evaluate(model, val_loader)
                curve_records.append({
                    "scheduler": scheduler_name, "seed": seed, "step": step,
                    "train_loss": None, "val_loss": vl, "val_acc": va, "lr": lr_now
                })
                if steps_to_target is None and vl <= TARGET_LOSS:
                    steps_to_target = step
                print(f"  [{scheduler_name}|seed={seed}] step {step}/{num_training_steps} "
                      f"val_loss={vl:.4f} val_acc={va:.4f} lr={lr_now:.2e}")

    final_val_loss, final_val_acc = evaluate(model, val_loader)
    elapsed = time.time() - start_time

    curve_records.append({
        "scheduler": scheduler_name, "seed": seed, "step": step,
        "train_loss": None, "val_loss": final_val_loss, "val_acc": final_val_acc, "lr": lr_now
    })

    val_points = [(r["step"], r["val_loss"]) for r in curve_records if r["val_loss"] is not None]
    val_steps = [p[0] for p in val_points]
    val_losses = [p[1] for p in val_points]
    auc = float(np.trapz(val_losses, val_steps)) if len(val_points) > 1 else float("inf")

    metrics = {
        "scheduler": scheduler_name,
        "seed": seed,
        "final_val_loss": final_val_loss,
        "final_val_acc": final_val_acc,
        "steps_to_target": steps_to_target,
        "auc_val_loss": auc,
        "total_steps": step,
        "elapsed_seconds": elapsed,
    }

    print(f"  [{scheduler_name}|seed={seed}] DONE — val_loss={final_val_loss:.4f} "
          f"val_acc={final_val_acc:.4f} steps_to_target={steps_to_target} "
          f"auc={auc:.1f} time={elapsed:.1f}s")

    del model, optimizer, scheduler
    torch.cuda.empty_cache() if torch.cuda.is_available() else None

    return metrics, curve_records


# ─── Aggregation ──────────────────────────────────────────────────────────────

def aggregate_results(all_per_run):
    df = pd.DataFrame(all_per_run)
    agg = df.groupby("scheduler").agg(
        final_val_loss_mean=("final_val_loss", "mean"),
        final_val_loss_std=("final_val_loss", "std"),
        final_val_acc_mean=("final_val_acc", "mean"),
        final_val_acc_std=("final_val_acc", "std"),
        auc_val_loss_mean=("auc_val_loss", "mean"),
        auc_val_loss_std=("auc_val_loss", "std"),
        steps_to_target_mean=("steps_to_target", lambda x: x.dropna().mean() if x.dropna().any() else None),
        steps_to_target_reached=("steps_to_target", lambda x: x.notna().sum()),
        elapsed_mean=("elapsed_seconds", "mean"),
        num_seeds=("seed", "count"),
    ).reset_index()
    agg = agg.sort_values("final_val_loss_mean")
    agg.insert(0, "rank", range(1, len(agg) + 1))
    return agg


# ─── Main ─────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="LR Scheduler Benchmark")
    parser.add_argument("--schedulers", nargs="+", default=list(SCHEDULER_REGISTRY.keys()),
                        choices=list(SCHEDULER_REGISTRY.keys()),
                        help="Which schedulers to benchmark")
    parser.add_argument("--seeds", type=int, default=3, help="Number of seeds per scheduler")
    parser.add_argument("--output-dir", type=str, default="results", help="Output directory")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Device: {DEVICE}")
    print(f"Model: {MODEL_NAME}")
    print(f"Dataset: {DATASET_NAME}/{DATASET_CONFIG}")
    print(f"Schedulers: {args.schedulers}")
    print(f"Seeds: {args.seeds}")
    print(f"Batch size: {BATCH_SIZE}, Epochs: {NUM_EPOCHS}, LR: {LEARNING_RATE}")
    print(f"Warmup: {WARMUP_FRACTION*100:.0f}%, Eval every: {EVAL_EVERY_STEPS} steps")
    print("=" * 80)

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    train_loader, val_loader = load_data(tokenizer)
    num_training_steps = len(train_loader) * NUM_EPOCHS
    print(f"Training steps per run: {num_training_steps}")
    print(f"Warmup steps: {int(WARMUP_FRACTION * num_training_steps)}")
    print("=" * 80)

    all_per_run = []
    all_curves = []

    for sched_name in args.schedulers:
        print(f"\n{'='*60}")
        print(f"Scheduler: {sched_name}")
        print(f"{'='*60}")
        for seed in range(args.seeds):
            print(f"\n--- Seed {seed} ---")
            metrics, curves = run_single(sched_name, seed, tokenizer, train_loader, val_loader)
            all_per_run.append(metrics)
            all_curves.extend(curves)

    per_run_df = pd.DataFrame(all_per_run)
    per_run_df.to_csv(output_dir / "per_run.csv", index=False)
    print(f"\nPer-run results saved to {output_dir / 'per_run.csv'}")

    curves_df = pd.DataFrame(all_curves)
    curves_df.to_csv(output_dir / "curves.csv", index=False)
    print(f"Curves saved to {output_dir / 'curves.csv'}")

    leaderboard_df = aggregate_results(all_per_run)
    leaderboard_df.to_csv(output_dir / "leaderboard.csv", index=False)
    print(f"Leaderboard saved to {output_dir / 'leaderboard.csv'}")

    config = {
        "model": MODEL_NAME,
        "dataset": f"{DATASET_NAME}/{DATASET_CONFIG}",
        "batch_size": BATCH_SIZE,
        "num_epochs": NUM_EPOCHS,
        "learning_rate": LEARNING_RATE,
        "weight_decay": WEIGHT_DECAY,
        "warmup_fraction": WARMUP_FRACTION,
        "eval_every_steps": EVAL_EVERY_STEPS,
        "target_loss": TARGET_LOSS,
        "num_seeds": args.seeds,
        "schedulers": args.schedulers,
        "device": DEVICE,
        "total_steps_per_run": num_training_steps,
    }
    with open(output_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    print("\n" + "=" * 80)
    print("LEADERBOARD")
    print("=" * 80)
    print(leaderboard_df.to_string(index=False))

    upload_results_to_space(output_dir)


def upload_results_to_space(output_dir):
    """Upload benchmark results to the HF Space data/ directory."""
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        space_id = "juiceb0xc0de/lr-scheduler-benchmark"

        for filename in ["leaderboard.csv", "curves.csv", "per_run.csv", "config.json"]:
            filepath = output_dir / filename
            if filepath.exists():
                api.upload_file(
                    path_or_fileobj=str(filepath),
                    path_in_repo=f"data/{filename}",
                    repo_id=space_id,
                    repo_type="space",
                )
                print(f"Uploaded {filename} to {space_id}/data/{filename}")

        print(f"\n✅ Results uploaded! Space: https://huggingface.co/spaces/{space_id}")
    except Exception as e:
        print(f"⚠️ Failed to upload to Space: {e}")
        print("Results are saved locally in the output directory.")


if __name__ == "__main__":
    main()