Njap-png
Add Agron neural-network zoo, optimisers, algorithms, schedulers and unified trainer
a311e94 | """Advanced optimisers for the Agron zoo. | |
| Every optimiser is a standard ``torch.optim.Optimizer`` subclass, so it can be | |
| used directly with PyTorch or routed through :func:`build_optimizer`. | |
| Registered optimisers | |
| --------------------- | |
| ``adam``, ``adamw``, ``adamax``, ``adagrad``, ``rmsprop``, ``sgd``, ``nadam``, | |
| ``lion``, ``adafactor``, ``adan``, ``lamb``, ``muon``, ``shampoo``, | |
| ``schedulefree``, ``kahan``, ``sophia`` | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple | |
| import torch | |
| from torch.optim import Optimizer | |
| # --------------------------------------------------------------------------- # | |
| # Registry | |
| # --------------------------------------------------------------------------- # | |
| OPTIMIZERS: Dict[str, Dict[str, Any]] = {} | |
| def _register(name: str, cls, description: str, family: str, defaults: Optional[Dict[str, Any]] = None) -> None: | |
| OPTIMIZERS[name] = { | |
| "name": name, | |
| "class": cls, | |
| "description": description, | |
| "family": family, | |
| "defaults": defaults or {}, | |
| } | |
| def optimizer_names() -> List[str]: | |
| return sorted(OPTIMIZERS.keys()) | |
| def get_optimizer(name: str) -> Dict[str, Any]: | |
| key = (name or "adamw").strip().lower() | |
| if key not in OPTIMIZERS: | |
| raise KeyError(f"unknown optimizer '{name}'. available: {', '.join(optimizer_names())}") | |
| return OPTIMIZERS[key] | |
| def build_optimizer(name: str, params, **overrides): | |
| """Instantiate the named optimiser for ``params``.""" | |
| entry = get_optimizer(name) | |
| defaults = dict(entry.get("defaults", {})) | |
| defaults.update(overrides) | |
| return entry["class"](params, **defaults) | |
| def optimizer_report() -> List[Dict[str, Any]]: | |
| return [ | |
| {"name": k, "family": v["family"], "description": v["description"], "defaults": v["defaults"]} | |
| for k, v in sorted(OPTIMIZERS.items()) | |
| ] | |
| def _check_lr(lr: float, eps: float = 0.0) -> None: | |
| if lr < 0.0: | |
| raise ValueError(f"Invalid learning rate: {lr}") | |
| if eps < 0.0: | |
| raise ValueError(f"Invalid epsilon value: {eps}") | |
| def _check_beta(beta: float, index: int) -> None: | |
| if not 0.0 <= beta < 1.0: | |
| raise ValueError(f"Invalid beta parameter at index {index}: {beta}") | |
| # --------------------------------------------------------------------------- # | |
| # Adam family | |
| # --------------------------------------------------------------------------- # | |
| class Adam(Optimizer): | |
| """Vanilla Adam with optional AMSGrad.""" | |
| def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999), | |
| eps: float = 1e-8, weight_decay: float = 0.0, amsgrad: bool = False): | |
| _check_lr(lr, eps) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("Adam does not support sparse gradients") | |
| if group["weight_decay"] != 0: | |
| grad = grad.add(p, alpha=group["weight_decay"]) | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["v"] = torch.zeros_like(p) | |
| if group["amsgrad"]: | |
| state["max_v"] = torch.zeros_like(p) | |
| state["step"] += 1 | |
| m, v = state["m"], state["v"] | |
| m.lerp_(grad, 1 - beta1) | |
| v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) | |
| bc1 = 1 - beta1 ** state["step"] | |
| bc2 = 1 - beta2 ** state["step"] | |
| if group["amsgrad"]: | |
| max_v = state["max_v"] | |
| max_v.copy_(torch.maximum(max_v, v)) | |
| denom = (max_v.sqrt() / math.sqrt(bc2)) + group["eps"] | |
| else: | |
| denom = (v.sqrt() / math.sqrt(bc2)) + group["eps"] | |
| p.addcdiv_(m, denom, value=-(group["lr"] / bc1)) | |
| return loss | |
| class AdamW(Optimizer): | |
| """AdamW with decoupled weight decay and optional AMSGrad.""" | |
| def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999), | |
| eps: float = 1e-8, weight_decay: float = 0.01, amsgrad: bool = False): | |
| _check_lr(lr, eps) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("AdamW does not support sparse gradients") | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["v"] = torch.zeros_like(p) | |
| if group["amsgrad"]: | |
| state["max_v"] = torch.zeros_like(p) | |
| state["step"] += 1 | |
| m, v = state["m"], state["v"] | |
| if group["weight_decay"] != 0: | |
| p.mul_(1 - group["lr"] * group["weight_decay"]) | |
| m.lerp_(grad, 1 - beta1) | |
| v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) | |
| bc1 = 1 - beta1 ** state["step"] | |
| bc2 = 1 - beta2 ** state["step"] | |
| if group["amsgrad"]: | |
| max_v = state["max_v"] | |
| max_v.copy_(torch.maximum(max_v, v)) | |
| denom = (max_v.sqrt() / math.sqrt(bc2)) + group["eps"] | |
| else: | |
| denom = (v.sqrt() / math.sqrt(bc2)) + group["eps"] | |
| p.addcdiv_(m, denom, value=-(group["lr"] / bc1)) | |
| return loss | |
| class NAdam(Optimizer): | |
| """Nesterov-accelerated Adam.""" | |
| def __init__(self, params, lr: float = 2e-3, betas: Tuple[float, float] = (0.9, 0.999), | |
| eps: float = 1e-8, weight_decay: float = 0.0, momentum_decay: float = 4e-3): | |
| _check_lr(lr, eps) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, | |
| momentum_decay=momentum_decay)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| mu = group["momentum_decay"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("NAdam does not support sparse gradients") | |
| if group["weight_decay"] != 0: | |
| grad = grad.add(p, alpha=group["weight_decay"]) | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["v"] = torch.zeros_like(p) | |
| state["mu_product"] = 1.0 | |
| state["step"] += 1 | |
| m, v = state["m"], state["v"] | |
| mu_product = state["mu_product"] * (1 - 0.5 * 0.96 ** (state["step"] * mu)) | |
| mu_product_next = mu_product * (1 - 0.5 * 0.96 ** ((state["step"] + 1) * mu)) | |
| state["mu_product"] = mu_product_next | |
| m.mul_(beta1).add_(grad, alpha=1 - beta1) | |
| v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) | |
| m_hat = beta1 * m / (1 - mu_product_next) + (1 - beta1) * grad / (1 - mu_product) | |
| v_hat = v / (1 - beta2 ** state["step"]) | |
| p.addcdiv_(m_hat, v_hat.sqrt().add_(group["eps"]), value=-group["lr"]) | |
| return loss | |
| class Adamax(Optimizer): | |
| """Adamax (infinity-norm variant of Adam).""" | |
| def __init__(self, params, lr: float = 2e-3, betas: Tuple[float, float] = (0.9, 0.999), | |
| eps: float = 1e-8, weight_decay: float = 0.0): | |
| _check_lr(lr, eps) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| if weight_decay < 0.0: | |
| raise ValueError(f"Invalid weight_decay value: {weight_decay}") | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("Adamax does not support sparse gradients") | |
| if group["weight_decay"] != 0: | |
| grad = grad.add(p, alpha=group["weight_decay"]) | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["u"] = torch.zeros_like(p) | |
| state["step"] += 1 | |
| m, u = state["m"], state["u"] | |
| m.mul_(beta1).add_(grad, alpha=1 - beta1) | |
| u.mul_(beta2).clamp_(min=0).add_(grad.abs().add_(group["eps"]), alpha=1 - beta2) | |
| bias_c = 1 - beta1 ** state["step"] | |
| p.addcdiv_(m, u, value=-(group["lr"] / bias_c)) | |
| return loss | |
| class Adagrad(Optimizer): | |
| """Adagrad with optional decoupled weight decay.""" | |
| def __init__(self, params, lr: float = 1e-2, lr_decay: float = 0.0, weight_decay: float = 0.0, | |
| initial_accumulator_value: float = 0.0, eps: float = 1e-10): | |
| _check_lr(lr, eps) | |
| if lr_decay < 0.0: | |
| raise ValueError(f"Invalid lr_decay value: {lr_decay}") | |
| if weight_decay < 0.0: | |
| raise ValueError(f"Invalid weight_decay value: {weight_decay}") | |
| if initial_accumulator_value < 0.0: | |
| raise ValueError("initial_accumulator_value must be non-negative") | |
| super().__init__(params, dict(lr=lr, lr_decay=lr_decay, weight_decay=weight_decay, | |
| initial_accumulator_value=initial_accumulator_value, eps=eps)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("Adagrad does not support sparse gradients") | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["sum"] = torch.full_like(p, group["initial_accumulator_value"]) | |
| state["step"] += 1 | |
| if group["weight_decay"] != 0: | |
| grad = grad.add(p, alpha=group["weight_decay"]) | |
| clr = group["lr"] / (1 + (state["step"] - 1) * group["lr_decay"]) | |
| state["sum"].addcmul_(grad, grad, value=1) | |
| p.addcdiv_(grad, state["sum"].sqrt().add_(group["eps"]), value=-clr) | |
| return loss | |
| class RMSprop(Optimizer): | |
| """RMSprop with momentum and centered variants.""" | |
| def __init__(self, params, lr: float = 1e-2, alpha: float = 0.99, eps: float = 1e-8, | |
| weight_decay: float = 0.0, momentum: float = 0.0, centered: bool = False): | |
| _check_lr(lr, eps) | |
| if momentum < 0.0: | |
| raise ValueError(f"Invalid momentum value: {momentum}") | |
| if weight_decay < 0.0: | |
| raise ValueError(f"Invalid weight_decay value: {weight_decay}") | |
| if not 0.0 <= alpha < 1.0: | |
| raise ValueError(f"Invalid alpha value: {alpha}") | |
| super().__init__(params, dict(lr=lr, alpha=alpha, eps=eps, weight_decay=weight_decay, | |
| momentum=momentum, centered=centered)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("RMSprop does not support sparse gradients") | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["square_avg"] = torch.zeros_like(p) | |
| if group["momentum"] > 0: | |
| state["momentum_buffer"] = torch.zeros_like(p) | |
| if group["centered"]: | |
| state["grad_avg"] = torch.zeros_like(p) | |
| state["step"] += 1 | |
| alpha = group["alpha"] | |
| square_avg = state["square_avg"] | |
| square_avg.mul_(alpha).addcmul_(grad, grad, value=1 - alpha) | |
| if group["centered"]: | |
| grad_avg = state["grad_avg"] | |
| grad_avg.mul_(alpha).add_(grad, alpha=1 - alpha) | |
| avg = square_avg.addcmul_(grad_avg, grad_avg, value=-1).sqrt().add_(group["eps"]) | |
| else: | |
| avg = square_avg.sqrt().add_(group["eps"]) | |
| if group["weight_decay"] != 0: | |
| grad = grad.add(p, alpha=group["weight_decay"]) | |
| if group["momentum"] > 0: | |
| buf = state["momentum_buffer"] | |
| buf.mul_(group["momentum"]).addcdiv_(grad, avg) | |
| p.add_(buf, alpha=-group["lr"]) | |
| else: | |
| p.addcdiv_(grad, avg, value=-group["lr"]) | |
| return loss | |
| class SGD(Optimizer): | |
| """SGD with momentum, dampening, weight decay and Nesterov.""" | |
| def __init__(self, params, lr: float = 1e-2, momentum: float = 0.0, dampening: float = 0.0, | |
| weight_decay: float = 0.0, nesterov: bool = False): | |
| _check_lr(lr) | |
| if momentum < 0.0: | |
| raise ValueError(f"Invalid momentum value: {momentum}") | |
| if weight_decay < 0.0: | |
| raise ValueError(f"Invalid weight_decay value: {weight_decay}") | |
| if nesterov and (momentum <= 0 or dampening != 0): | |
| raise ValueError("Nesterov momentum requires momentum > 0 and dampening == 0") | |
| super().__init__(params, dict(lr=lr, momentum=momentum, dampening=dampening, | |
| weight_decay=weight_decay, nesterov=nesterov)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| d_p = p.grad | |
| if group["weight_decay"] != 0: | |
| d_p = d_p.add(p, alpha=group["weight_decay"]) | |
| if group["momentum"] != 0: | |
| buf = self.state[p].get("momentum_buffer") | |
| if buf is None: | |
| buf = self.state[p]["momentum_buffer"] = torch.clone(d_p).detach() | |
| else: | |
| buf.mul_(group["momentum"]).add_(d_p, alpha=1 - group["dampening"]) | |
| d_p = d_p.add(buf, alpha=group["momentum"]) if group["nesterov"] else buf | |
| p.add_(d_p, alpha=-group["lr"]) | |
| return loss | |
| # --------------------------------------------------------------------------- # | |
| # Sign / momentum based | |
| # --------------------------------------------------------------------------- # | |
| class Lion(Optimizer): | |
| """Evolved Sign Momentum (Chen et al., 2023).""" | |
| def __init__(self, params, lr: float = 1e-4, betas: Tuple[float, float] = (0.9, 0.99), | |
| weight_decay: float = 0.01): | |
| _check_lr(lr) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| super().__init__(params, dict(lr=lr, betas=betas, weight_decay=weight_decay)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("Lion does not support sparse gradients") | |
| state = self.state[p] | |
| if not state: | |
| state["m"] = torch.zeros_like(p) | |
| m = state["m"] | |
| if group["weight_decay"] != 0: | |
| p.mul_(1 - group["lr"] * group["weight_decay"]) | |
| update = m.mul(beta1).add(grad, alpha=1 - beta1) | |
| p.add_(torch.sign(update), alpha=-group["lr"]) | |
| m.mul_(beta2).add_(grad, alpha=1 - beta2) | |
| state["m"] = m | |
| return loss | |
| # --------------------------------------------------------------------------- # | |
| # Factored / memory efficient | |
| # --------------------------------------------------------------------------- # | |
| class Adafactor(Optimizer): | |
| """Adafactor with factored second moments and update clipping.""" | |
| def __init__(self, params, lr: float = 1e-3, beta2: float = 0.999, | |
| eps: Tuple[float, float] = (1e-30, 1e-3), clip_threshold: float = 1.0, | |
| weight_decay: float = 0.0, min_dim_size_to_factor: int = 128): | |
| _check_lr(lr) | |
| _check_beta(beta2, 1) | |
| super().__init__(params, dict(lr=lr, beta2=beta2, eps=eps, clip_threshold=clip_threshold, | |
| weight_decay=weight_decay, | |
| min_dim_size_to_factor=min_dim_size_to_factor)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta2 = group["beta2"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("Adafactor does not support sparse gradients") | |
| state = self.state[p] | |
| factored = ( | |
| grad.dim() >= 2 | |
| and min(grad.shape[-1], grad.shape[-2]) >= group["min_dim_size_to_factor"] | |
| ) | |
| if not state: | |
| state["step"] = 0 | |
| if factored: | |
| state["row"] = torch.zeros(grad.shape[:-1], device=grad.device) | |
| state["col"] = torch.zeros(grad.shape[:-2] + grad.shape[-1:], device=grad.device) | |
| else: | |
| state["v"] = torch.zeros_like(grad) | |
| state["step"] += 1 | |
| step = state["step"] | |
| if factored: | |
| row, col = state["row"], state["col"] | |
| row.mul_(beta2).add_(grad.pow(2).mean(dim=-1), alpha=1 - beta2) | |
| col.mul_(beta2).add_(grad.pow(2).mean(dim=-2), alpha=1 - beta2) | |
| row_mean = row.mean(dim=-1, keepdim=True).clamp_min(1e-30) | |
| col_mean = col.mean(dim=-2, keepdim=True).clamp_min(1e-30) | |
| v = (row / row_mean).unsqueeze(-1) * (col / col_mean).unsqueeze(-2) | |
| else: | |
| v = state["v"] | |
| v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) | |
| update = grad / (v.sqrt() + group["eps"][1]) | |
| update_norm = update.norm() | |
| if update_norm > 0: | |
| update = update / max(1.0, (update_norm / group["clip_threshold"]).item()) | |
| rel_step = group["lr"] * math.sqrt(max(1e-38, 1 - beta2 ** step)) | |
| p.add_(update, alpha=-rel_step) | |
| if group["weight_decay"] != 0: | |
| p.mul_(1 - group["lr"] * group["weight_decay"]) | |
| return loss | |
| class Adan(Optimizer): | |
| """Adan - adaptive Nesterov momentum optimiser (Xie et al., 2022).""" | |
| def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float, float] = (0.98, 0.92, 0.99), | |
| eps: float = 1e-8, weight_decay: float = 0.01): | |
| _check_lr(lr, eps) | |
| for i, b in enumerate(betas): | |
| _check_beta(b, i) | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2, beta3 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("Adan does not support sparse gradients") | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["v"] = torch.zeros_like(p) | |
| state["n"] = torch.zeros_like(p) | |
| state["prev_grad"] = torch.zeros_like(p) | |
| state["step"] += 1 | |
| m, v, n, prev_grad = state["m"], state["v"], state["n"], state["prev_grad"] | |
| if group["weight_decay"] != 0: | |
| p.mul_(1 - group["lr"] * group["weight_decay"]) | |
| grad_diff = grad - prev_grad | |
| m.mul_(beta1).add_(grad, alpha=1 - beta1) | |
| v.mul_(beta2).add_(grad_diff, alpha=1 - beta2) | |
| n.mul_(beta3).addcmul_(grad + (1 - beta2) * grad_diff, | |
| grad + (1 - beta2) * grad_diff, value=1 - beta3) | |
| prev_grad.copy_(grad) | |
| bc1 = 1 - beta1 ** state["step"] | |
| bc2 = 1 - beta2 ** state["step"] | |
| bc3 = 1 - beta3 ** state["step"] | |
| denom = (n.sqrt() / math.sqrt(bc3)) + group["eps"] | |
| update = (m + (1 - beta2) * v) / denom | |
| p.add_(update, alpha=-(group["lr"] / bc1)) | |
| return loss | |
| class LAMB(Optimizer): | |
| """LAMB - layer-wise adaptive moments for batched training.""" | |
| def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999), | |
| eps: float = 1e-6, weight_decay: float = 0.01): | |
| _check_lr(lr, eps) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("LAMB does not support sparse gradients") | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["v"] = torch.zeros_like(p) | |
| state["step"] += 1 | |
| m, v = state["m"], state["v"] | |
| m.mul_(beta1).add_(grad, alpha=1 - beta1) | |
| v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) | |
| bc1 = 1 - beta1 ** state["step"] | |
| bc2 = 1 - beta2 ** state["step"] | |
| update = (m / bc1) / ((v / bc2).sqrt() + group["eps"]) | |
| if group["weight_decay"] != 0: | |
| update = update.add(p, alpha=group["weight_decay"]) | |
| param_norm = p.norm() | |
| update_norm = update.norm() | |
| trust_ratio = 1.0 | |
| if param_norm > 0 and update_norm > 0: | |
| trust_ratio = (param_norm / update_norm).item() | |
| p.add_(update, alpha=-group["lr"] * trust_ratio) | |
| return loss | |
| # --------------------------------------------------------------------------- # | |
| # Second-order inspired | |
| # --------------------------------------------------------------------------- # | |
| def _zeropower_newton_schulz(matrix: torch.Tensor, steps: int = 5) -> torch.Tensor: | |
| """Approximate the orthogonal polar factor via Newton-Schulz iteration.""" | |
| a, b, c = 3.4445, -4.7750, 2.0315 | |
| x = matrix.float() | |
| transposed = x.size(0) > x.size(1) | |
| if transposed: | |
| x = x.T | |
| x = x / (x.norm() + 1e-7) | |
| for _ in range(steps): | |
| gram = x @ x.T | |
| x = a * x + (b * gram + c * (gram @ gram)) @ x | |
| if transposed: | |
| x = x.T | |
| return x.to(matrix.dtype) | |
| class Muon(Optimizer): | |
| """Muon - momentum orthogonalised by Newton-Schulz iterations. | |
| Two-dimensional parameters are orthogonalised; one-dimensional parameters | |
| fall back to Adam-style updates. | |
| """ | |
| def __init__(self, params, lr: float = 0.02, momentum: float = 0.95, | |
| nesterov: bool = True, weight_decay: float = 0.0, | |
| ns_steps: int = 5, eps: float = 1e-8, | |
| adam_betas: Tuple[float, float] = (0.9, 0.999)): | |
| _check_lr(lr, eps) | |
| super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, | |
| weight_decay=weight_decay, ns_steps=ns_steps, | |
| eps=eps, adam_betas=adam_betas, | |
| ns_coefficients=(3.4445, -4.7750, 2.0315))) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| params = [p for p in group["params"] if p.grad is not None] | |
| if not params: | |
| continue | |
| grads = [p.grad for p in params] | |
| if group["weight_decay"] != 0: | |
| for p in params: | |
| p.mul_(1 - group["lr"] * group["weight_decay"]) | |
| # 2D parameters: orthogonalised momentum | |
| ortho = [(p, g) for p, g in zip(params, grads) if p.dim() >= 2] | |
| flat = [(p, g) for p, g in zip(params, grads) if p.dim() < 2] | |
| for p, g in ortho: | |
| state = self.state[p] | |
| if "momentum" not in state: | |
| state["momentum"] = torch.zeros_like(p) | |
| buf = state["momentum"] | |
| buf.mul_(group["momentum"]).add_(g) | |
| update = g.add(buf, alpha=group["momentum"]) if group["nesterov"] else buf | |
| o = _zeropower_newton_schulz(update, group["ns_steps"]) | |
| scale = max(1.0, p.size(0) / p.size(1)) ** 0.5 | |
| p.add_(o, alpha=-group["lr"] * scale) | |
| beta1, beta2 = group["adam_betas"] | |
| for p, g in flat: | |
| state = self.state[p] | |
| if "step" not in state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["v"] = torch.zeros_like(p) | |
| state["step"] += 1 | |
| m, v = state["m"], state["v"] | |
| m.lerp_(g, 1 - beta1) | |
| v.mul_(beta2).addcmul_(g, g, value=1 - beta2) | |
| bc1 = 1 - beta1 ** state["step"] | |
| bc2 = 1 - beta2 ** state["step"] | |
| denom = (v.sqrt() / math.sqrt(bc2)) + group["eps"] | |
| p.addcdiv_(m, denom, value=-(group["lr"] / bc1)) | |
| return loss | |
| def _matrix_inverse_root(mat: torch.Tensor, root: int = 2, eps: float = 1e-8) -> torch.Tensor: | |
| """``mat ** (-1/root)`` via a symmetric eigendecomposition.""" | |
| mat = mat.float() + eps * torch.eye(mat.size(0), device=mat.device) | |
| eigvals, eigvecs = torch.linalg.eigh(mat) | |
| eigvals = eigvals.clamp_min(1e-12).pow(-1.0 / root) | |
| return (eigvecs * eigvals.unsqueeze(0)) @ eigvecs.T | |
| class Shampoo(Optimizer): | |
| """Shampoo - full-matrix preconditioning of gradients. | |
| Preconditioners are refreshed every ``precondition_frequency`` steps and | |
| fall back to a diagonal preconditioner for very large dimensions. | |
| """ | |
| def __init__(self, params, lr: float = 1e-2, betas: Tuple[float, float] = (0.9, 0.999), | |
| eps: float = 1e-8, weight_decay: float = 0.0, | |
| precondition_frequency: int = 25, max_dim: int = 2048): | |
| _check_lr(lr, eps) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, | |
| precondition_frequency=precondition_frequency, | |
| max_dim=max_dim)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if grad.is_sparse: | |
| raise RuntimeError("Shampoo does not support sparse gradients") | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["precond"] = None | |
| state["factors"] = {} | |
| state["step"] += 1 | |
| m = state["m"] | |
| m.mul_(beta1).add_(grad, alpha=1 - beta1) | |
| update = m | |
| if group["weight_decay"] != 0: | |
| p.mul_(1 - group["lr"] * group["weight_decay"]) | |
| if grad.dim() == 2: | |
| factors = state["factors"] | |
| if "L" not in factors: | |
| factors["L"] = torch.zeros_like(grad @ grad.T) | |
| factors["R"] = torch.zeros_like(grad.T @ grad) | |
| factors["L"].mul_(beta2).add_(grad @ grad.T, alpha=1 - beta2) | |
| factors["R"].mul_(beta2).add_(grad.T @ grad, alpha=1 - beta2) | |
| if state["step"] % group["precondition_frequency"] == 0 or state["precond"] is None: | |
| if max(grad.shape) > group["max_dim"]: | |
| # diagonal fallback for very large layers | |
| diagonal = factors["L"].diagonal().clamp_min(1e-12).pow(-0.25) | |
| l_inv = torch.diag(diagonal) | |
| r_diag = factors["R"].diagonal().clamp_min(1e-12).pow(-0.25) | |
| r_inv = torch.diag(r_diag) | |
| else: | |
| l_inv = _matrix_inverse_root(factors["L"], 4) | |
| r_inv = _matrix_inverse_root(factors["R"], 4) | |
| state["precond"] = (l_inv, r_inv) | |
| l_inv, r_inv = state["precond"] | |
| try: | |
| update = l_inv @ update @ r_inv | |
| except Exception: | |
| update = update | |
| elif grad.dim() > 2: | |
| update = update / (grad.pow(2).mean(dim=tuple(range(1, grad.dim())), keepdim=True).sqrt() + group["eps"]) | |
| p.add_(update, alpha=-group["lr"]) | |
| return loss | |
| class Sophia(Optimizer): | |
| """Sophia - second-order clipped stochastic optimisation. | |
| The diagonal Hessian estimate is updated by :meth:`update_hessian` using | |
| the Gauss-Newton-Bartlett (GNB) estimator with a Hutchinson probe. | |
| """ | |
| def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float] = (0.965, 0.99), | |
| eps: float = 1e-15, weight_decay: float = 0.01, rho: float = 0.04): | |
| _check_lr(lr, eps) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, rho=rho)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["h"] = torch.zeros_like(p) | |
| state["step"] += 1 | |
| m, h = state["m"], state["h"] | |
| if group["weight_decay"] != 0: | |
| p.mul_(1 - group["lr"] * group["weight_decay"]) | |
| m.mul_(beta1).add_(grad, alpha=1 - beta1) | |
| ratio = (m / (group["rho"] * h + group["eps"])).clamp(-1.0, 1.0) | |
| p.add_(ratio, alpha=-group["lr"]) | |
| return loss | |
| def update_hessian(self, loss: torch.Tensor, seed: Optional[int] = None) -> None: | |
| """Refresh the diagonal Hessian using a Hutchinson/GNB probe.""" | |
| params = [p for group in self.param_groups for p in group["params"] if p.requires_grad] | |
| if not params: | |
| return | |
| generator = torch.Generator(device=params[0].device) | |
| if seed is not None: | |
| generator.manual_seed(seed) | |
| grads = torch.autograd.grad(loss, params, create_graph=True, retain_graph=True, allow_unused=True) | |
| probes = [torch.randint_like(p, 0, 2, generator=generator).float() * 2 - 1 for p in params] | |
| dot = sum((g * v).sum() for g, v in zip(grads, probes) if g is not None and g.requires_grad) | |
| if not isinstance(dot, torch.Tensor) or not dot.requires_grad: | |
| return | |
| second = torch.autograd.grad(dot, params, retain_graph=False, allow_unused=True) | |
| for group in self.param_groups: | |
| beta2 = group["betas"][1] | |
| for p, hv in zip(params, second): | |
| if p not in self.state: | |
| continue | |
| if hv is None: | |
| continue | |
| state = self.state[p] | |
| h = state["h"] | |
| probe_sq = (hv * hv).abs() | |
| h.mul_(beta2).add_(probe_sq, alpha=1 - beta2) | |
| # --------------------------------------------------------------------------- # | |
| # Schedule-free and compensated | |
| # --------------------------------------------------------------------------- # | |
| class ScheduleFreeAdamW(Optimizer): | |
| """Schedule-Free AdamW (Defazio et al., 2024). | |
| Implements Algorithm 1 of the paper: ``z`` is the fast iterate, ``x`` the | |
| averaged eval point and ``y = (1-b1) z + b1 x`` the gradient point. | |
| Call :meth:`train` before the forward pass and :meth:`eval` before | |
| validation / inference. | |
| """ | |
| def __init__(self, params, lr: float = 2.5e-3, betas: Tuple[float, float] = (0.9, 0.999), | |
| eps: float = 1e-8, weight_decay: float = 0.0, warmup_steps: int = 0): | |
| _check_lr(lr, eps) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, | |
| warmup_steps=warmup_steps)) | |
| def _state(self, p): | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["z"] = p.detach().clone() | |
| state["x"] = p.detach().clone() | |
| state["v"] = torch.zeros_like(p) | |
| return state | |
| def train(self): | |
| """Set the parameters to the gradient point ``y``.""" | |
| for group in self.param_groups: | |
| beta1 = group["betas"][0] | |
| for p in group["params"]: | |
| state = self._state(p) | |
| p.copy_((1 - beta1) * state["z"] + beta1 * state["x"]) | |
| def eval(self): | |
| """Set the parameters to the averaged eval point ``x``.""" | |
| for group in self.param_groups: | |
| for p in group["params"]: | |
| state = self.state[p] | |
| if "x" in state: | |
| p.copy_(state["x"]) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| lr = group["lr"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| state = self._state(p) | |
| state["step"] += 1 | |
| t = state["step"] | |
| z, x, v = state["z"], state["x"], state["v"] | |
| grad = p.grad | |
| v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) | |
| if t == 1: | |
| v.copy_(grad.pow(2)) | |
| denom = v.sqrt().add_(group["eps"]) | |
| if group["weight_decay"] != 0: | |
| z.mul_(1 - lr * group["weight_decay"]) | |
| z.addcdiv_(grad, denom, value=-lr) | |
| warmup = group["warmup_steps"] | |
| ckp1 = t / warmup if (warmup > 0 and t < warmup) else 1.0 / (t + 1) | |
| x.lerp_(z, ckp1) | |
| # prepare the next gradient point | |
| p.copy_((1 - beta1) * z + beta1 * x) | |
| return loss | |
| class KahanAdamW(Optimizer): | |
| """AdamW with Kahan compensated summation on the parameter update. | |
| Compensated summation removes the low-precision drift that accumulates | |
| when training long runs in bf16/fp16. | |
| """ | |
| def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999), | |
| eps: float = 1e-8, weight_decay: float = 0.01): | |
| _check_lr(lr, eps) | |
| _check_beta(betas[0], 0) | |
| _check_beta(betas[1], 1) | |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)) | |
| def step(self, closure=None): | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() | |
| for group in self.param_groups: | |
| beta1, beta2 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad | |
| if group["weight_decay"] != 0: | |
| grad = grad.add(p, alpha=group["weight_decay"]) | |
| state = self.state[p] | |
| if not state: | |
| state["step"] = 0 | |
| state["m"] = torch.zeros_like(p) | |
| state["v"] = torch.zeros_like(p) | |
| state["comp"] = torch.zeros_like(p, dtype=torch.float32) | |
| state["step"] += 1 | |
| m, v = state["m"], state["v"] | |
| m.lerp_(grad, 1 - beta1) | |
| v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) | |
| bc1 = 1 - beta1 ** state["step"] | |
| bc2 = 1 - beta2 ** state["step"] | |
| denom = (v.sqrt() / math.sqrt(bc2)) + group["eps"] | |
| delta = -(group["lr"] / bc1) * (m / denom) | |
| # Kahan compensated update | |
| comp = state["comp"] | |
| y = delta.float() - comp | |
| t = p.float() + y | |
| comp = (t - p.float()) - y | |
| state["comp"] = comp | |
| p.copy_(t.to(p.dtype)) | |
| return loss | |
| # --------------------------------------------------------------------------- # | |
| # Register everything | |
| # --------------------------------------------------------------------------- # | |
| _register("adam", Adam, "Vanilla Adam with optional AMSGrad", "adaptive", | |
| {"lr": 1e-3, "betas": (0.9, 0.999)}) | |
| _register("adamw", AdamW, "AdamW with decoupled weight decay", "adaptive", | |
| {"lr": 1e-3, "betas": (0.9, 0.999), "weight_decay": 0.01}) | |
| _register("nadam", NAdam, "Nesterov-accelerated Adam", "adaptive", | |
| {"lr": 2e-3, "betas": (0.9, 0.999)}) | |
| _register("adamax", Adamax, "Adam with an infinity-norm second moment", "adaptive", | |
| {"lr": 2e-3, "betas": (0.9, 0.999)}) | |
| _register("adagrad", Adagrad, "Cumulative squared-gradient preconditioning", "adaptive", | |
| {"lr": 1e-2}) | |
| _register("rmsprop", RMSprop, "RMSprop with momentum / centering", "adaptive", | |
| {"lr": 1e-2, "alpha": 0.99}) | |
| _register("sgd", SGD, "SGD with momentum and Nesterov", "momentum", | |
| {"lr": 1e-2, "momentum": 0.9}) | |
| _register("lion", Lion, "Evolved Sign Momentum (Lion)", "sign-momentum", | |
| {"lr": 1e-4, "betas": (0.9, 0.99), "weight_decay": 0.01}) | |
| _register("adafactor", Adafactor, "Memory-efficient factored second moments", "factored", | |
| {"lr": 1e-3}) | |
| _register("adan", Adan, "Adaptive Nesterov momentum (Adan)", "adaptive", | |
| {"lr": 1e-3, "betas": (0.98, 0.92, 0.99), "weight_decay": 0.01}) | |
| _register("lamb", LAMB, "Layer-wise adaptive moments (LAMB)", "layer-adaptive", | |
| {"lr": 1e-3, "betas": (0.9, 0.999), "weight_decay": 0.01}) | |
| _register("muon", Muon, "Momentum orthogonalised with Newton-Schulz", "orthogonalised", | |
| {"lr": 0.02, "momentum": 0.95}) | |
| _register("shampoo", Shampoo, "Full-matrix gradient preconditioning", "second-order", | |
| {"lr": 1e-2, "precondition_frequency": 25}) | |
| _register("sophia", Sophia, "Second-order clipped optimisation with GNB Hessian", "second-order", | |
| {"lr": 1e-3, "betas": (0.965, 0.99)}) | |
| _register("schedulefree", ScheduleFreeAdamW, "Schedule-Free AdamW", "schedule-free", | |
| {"lr": 2.5e-3, "betas": (0.9, 0.999)}) | |
| _register("kahan", KahanAdamW, "AdamW with Kahan compensated updates", "compensated", | |
| {"lr": 1e-3, "betas": (0.9, 0.999), "weight_decay": 0.01}) | |
| __all__ = [ | |
| "OPTIMIZERS", | |
| "Adam", | |
| "AdamW", | |
| "NAdam", | |
| "Adamax", | |
| "Adagrad", | |
| "RMSprop", | |
| "SGD", | |
| "Lion", | |
| "Adafactor", | |
| "Adan", | |
| "LAMB", | |
| "Muon", | |
| "Shampoo", | |
| "Sophia", | |
| "ScheduleFreeAdamW", | |
| "KahanAdamW", | |
| "build_optimizer", | |
| "get_optimizer", | |
| "optimizer_names", | |
| "optimizer_report", | |
| ] | |