""" 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()