Merge pull request #366 from kozistr/feature/optimizers

[Feature] AdaGC and SimplifiedAdEMAMix optimizers

Author: kozistr

README.md

@@ -5,6 +5,10 @@
 * Support `StableSPAM` optimizer. (#358, #359)
     * [How to Train in 4-Bit More Stably than 16-Bit Adam](https://arxiv.org/abs/2502.17055?)
 * Support `ScheduleFreeWrapper`. (#334, #360)
+* Implement `AdaGC` optimizer. (#364, #366)
+    * [Improving Training Stability for Large Language Model Pretraining](https://arxiv.org/abs/2502.11034)
+* Implement `Simplified-Ademamix` optimizer. (#364, #366)
+    * [Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants](https://arxiv.org/abs/2502.02431)
 
 ### Update
 

docs/changelogs/v3.4.3.md renamed to docs/changelogs/v3.5.0.md

@@ -28,6 +28,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.AdaGC
+    :docstring:
+    :members:
+
 ::: pytorch_optimizer.AdaHessian
     :docstring:
     :members:
@@ -92,6 +96,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.SimplifiedAdEMAMix
+    :docstring:
+    :members:
+
 ::: pytorch_optimizer.ADOPT
     :docstring:
     :members:

docs/index.md

@@ -22,6 +22,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_AdaFactor.png)
 
+### AdaGC
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_AdaGC.png)
+
 ### AdaHessian
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_AdaHessian.png)
@@ -326,6 +330,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_SignSGD.png)
 
+### SimplifiedAdEMAMix
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_SimplifiedAdEMAMix.png)
+
 ### SM3
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_SM3.png)
@@ -392,6 +400,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_AdaFactor.png)
 
+### AdaGC
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_AdaGC.png)
+
 ### AdaHessian
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_AdaHessian.png)
@@ -696,6 +708,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_SignSGD.png)
 
+### SimplifiedAdEMAMix
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_SimplifiedAdEMAMix.png)
+
 ### SM3
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_SM3.png)

docs/optimizer.md

@@ -11,17 +11,17 @@ repository = "https://github.com/kozistr/pytorch_optimizer"
 documentation = "https://pytorch-optimizers.readthedocs.io/en/latest"
 keywords = [
     "pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "ASGD", "AccSGD", "AdaBelief", "AdaBound",
-    "AdaDelta", "AdaFactor", "AdaMax", "AdamG", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdEMAMix", "ADOPT",
-    "AdaHessian", "Adai", "Adalite", "AdaLomo", "AdamMini", "AdamP", "AdamS", "Adan", "AggMo", "Aida", "AliG", "Amos",
-    "Apollo", "APOLLO", "AvaGrad", "bSAM", "CAME", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD",
-    "DAdaptLion", "DeMo", "DiffGrad", "EXAdam", "FAdam", "FOCUS", "Fromage", "FTRL", "GaLore", "Grams", "Gravity",
-    "GrokFast", "GSAM", "Kate", "Lamb", "LaProp", "LARS", "Lion", "LOMO", "Lookahead", "MADGRAD", "MARS", "MSVAG",
-    "Muno", "Nero", "NovoGrad", "OrthoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "PSGD", "QHAdam", "QHM",
-    "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM", "GCSAM", "LookSAM", "ScheduleFreeSGD", "ScheduleFreeAdamW",
-    "ScheduleFreeRAdam", "SCION", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH",
-    "SPAM", "StableSPAM", "SRMM", "StableAdamW", "SWATS", "TAM", "Tiger", "TRAC", "WSAM", "Yogi", "BCE", "BCEFocal",
-    "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky",
-    "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
+    "AdaDelta", "AdaFactor", "AdaGC", "AdaMax", "AdamG", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdEMAMix",
+    "Simplified-AdEMAMix", "ADOPT", "AdaHessian", "Adai", "Adalite", "AdaLomo", "AdamMini", "AdamP", "AdamS", "Adan",
+    "AggMo", "Aida", "AliG", "Amos", "Apollo", "APOLLO", "AvaGrad", "bSAM", "CAME", "DAdaptAdaGrad", "DAdaptAdam",
+    "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DeMo", "DiffGrad", "EXAdam", "FAdam", "FOCUS", "Fromage", "FTRL",
+    "GaLore", "Grams", "Gravity", "GrokFast", "GSAM", "Kate", "Lamb", "LaProp", "LARS", "Lion", "LOMO", "Lookahead",
+    "MADGRAD", "MARS", "MSVAG", "Muno", "Nero", "NovoGrad", "OrthoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy",
+    "PSGD", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM", "GCSAM", "LookSAM", "ScheduleFreeSGD",
+    "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SCION", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3",
+    "SOAP", "SopihaH", "SPAM", "StableSPAM", "SRMM", "StableAdamW", "SWATS", "TAM", "Tiger", "TRAC", "WSAM", "Yogi",
+    "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky",
+    "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",

docs/visualization.md

@@ -75,6 +75,7 @@
     AdaBound,
     AdaDelta,
     AdaFactor,
+    AdaGC,
     AdaHessian,
     Adai,
     Adalite,
@@ -143,6 +144,7 @@
     SGDSaI,
     Shampoo,
     SignSGD,
+    SimplifiedAdEMAMix,
     SophiaH,
     StableAdamW,
     StableSPAM,

docs/visualizations/rastrigin_AdaGC.png

@@ -13,6 +13,7 @@
 from pytorch_optimizer.optimizer.adabound import AdaBound
 from pytorch_optimizer.optimizer.adadelta import AdaDelta
 from pytorch_optimizer.optimizer.adafactor import AdaFactor
+from pytorch_optimizer.optimizer.adagc import AdaGC
 from pytorch_optimizer.optimizer.adahessian import AdaHessian
 from pytorch_optimizer.optimizer.adai import Adai
 from pytorch_optimizer.optimizer.adalite import Adalite
@@ -28,7 +29,7 @@
 from pytorch_optimizer.optimizer.adapnm import AdaPNM
 from pytorch_optimizer.optimizer.adashift import AdaShift
 from pytorch_optimizer.optimizer.adasmooth import AdaSmooth
-from pytorch_optimizer.optimizer.ademamix import AdEMAMix
+from pytorch_optimizer.optimizer.ademamix import AdEMAMix, SimplifiedAdEMAMix
 from pytorch_optimizer.optimizer.adopt import ADOPT
 from pytorch_optimizer.optimizer.agc import agc
 from pytorch_optimizer.optimizer.aggmo import AggMo
@@ -292,6 +293,7 @@ def load_optimizer(optimizer: str) -> OPTIMIZER:
     AdaLOMO,
     AdamG,
     AdEMAMix,
+    SimplifiedAdEMAMix,
     SOAP,
     ADOPT,
     FTRL,
@@ -308,6 +310,7 @@ def load_optimizer(optimizer: str) -> OPTIMIZER:
     EXAdam,
     SCION,
     StableSPAM,
+    AdaGC,
     Ranger25,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}

docs/visualizations/rastrigin_SimplifiedAdEMAMix.png

@@ -0,0 +1,139 @@
+import math
+
+import torch
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.type import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+from pytorch_optimizer.optimizer.utils import get_global_gradient_norm
+
+
+class AdaGC(BaseOptimizer):
+    r"""Improving Training Stability for Large Language Model Pretraining.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace.
+    :param beta: float. smoothing coefficient for EMA.
+    :param lambda_abs: float. absolute clipping threshold to prevent unstable updates from gradient explosions.
+    :param lambda_rel: float. relative clipping threshold to prevent unstable updates from gradient explosions.
+    :param warmup_steps: int. warmup steps.
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    :param fixed_decay: bool. fix weight decay.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        betas: BETAS = (0.9, 0.999),
+        beta: float = 0.98,
+        lambda_abs: float = 1.0,
+        lambda_rel: float = 1.05,
+        warmup_steps: int = 100,
+        weight_decay: float = 1e-1,
+        weight_decouple: bool = True,
+        fixed_decay: bool = False,
+        eps: float = 1e-8,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_range(beta, 'beta', 0.0, 1.0, '[)')
+        self.validate_positive(lambda_abs, 'lambda_abs')
+        self.validate_positive(lambda_rel, 'lambda_rel')
+        self.validate_non_negative(warmup_steps, 'warmup_steps')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'beta': beta,
+            'lambda_abs': lambda_abs,
+            'lambda_rel': lambda_rel,
+            'warmup_steps': warmup_steps,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'eps': eps,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'AdaGC'
+
+    @torch.no_grad()
+    def reset(self):
+        pass
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            beta1, beta2 = group['betas']
+
+            bias_correction1: float = self.debias(beta1, group['step'])
+            bias_correction2_sq: float = math.sqrt(self.debias(beta2, group['step']))
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(str(self))
+
+                state = self.state[p]
+
+                if 'exp_avg' not in state:
+                    state['exp_avg'] = torch.zeros_like(grad)
+                    state['exp_avg_sq'] = torch.zeros_like(grad)
+                    state['gamma'] = torch.empty((1,), device=grad.device, dtype=grad.dtype)
+
+                self.apply_weight_decay(
+                    p=p,
+                    grad=grad,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=group['fixed_decay'],
+                )
+
+                gamma = state['gamma']
+
+                if group['step'] < group['warmup_steps']:
+                    grad_norm = get_global_gradient_norm(self.param_groups).add_(group['eps'])
+
+                    h_t = min(group['lambda_abs'] / grad_norm, 1.0)
+                    g_hat = grad.mul(h_t)
+
+                    g_hat_norm = g_hat.norm()
+
+                    gamma.copy_(g_hat_norm if group['step'] == 1 else min(gamma, g_hat_norm))
+                else:
+                    h_t = min(group['lambda_rel'] * gamma / grad.norm(), 1.0)
+                    g_hat = grad.mul(h_t)
+
+                    gamma.mul_(group['beta']).add_(g_hat.norm(), alpha=1.0 - group['beta'])
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                exp_avg.mul_(beta1).add_(g_hat, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(g_hat, g_hat, value=1.0 - beta2)
+
+                update = (exp_avg / bias_correction1) / exp_avg_sq.sqrt().div_(bias_correction2_sq).add_(group['eps'])
+
+                p.add_(update, alpha=-group['lr'])
+
+        return loss