pytorch-optimizer 3.3.2__tar.gz → 3.3.3__tar.gz

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/PKG-INFO

@@ -1,10 +1,9 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: pytorch_optimizer
-Version: 3.3.2
+Version: 3.3.3
 Summary: optimizer & lr scheduler & objective function collections in PyTorch
-Home-page: https://github.com/kozistr/pytorch_optimizer
 License: Apache-2.0
-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,FAdam,Fromage,FTRL,GaLore,Gravity,GrokFast,GSAM,Kate,Lamb,LaProp,LARS,Lion,LOMO,Lookahead,MADGRAD,MARS,MSVAG,Muno,Nero,NovoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,ScheduleFreeRAdam,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SOAP,SopihaH,SRMM,StableAdamW,SWATS,Tiger,TRAC,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD,QGaLore
+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,FAdam,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,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,ScheduleFreeRAdam,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SOAP,SopihaH,SRMM,StableAdamW,SWATS,Tiger,TRAC,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD,QGaLore
 Author: kozistr
 Author-email: kozistr@gmail.com
 Maintainer: kozistr
@@ -36,6 +35,7 @@ Requires-Dist: numpy (<=1.24.4) ; python_version < "3.9"
 Requires-Dist: numpy (>1.24.4) ; python_version >= "3.9"
 Requires-Dist: torch (>=1.10) ; python_version >= "3.8"
 Project-URL: Documentation, https://pytorch-optimizers.readthedocs.io/en/latest
+Project-URL: Homepage, https://github.com/kozistr/pytorch_optimizer
 Project-URL: Repository, https://github.com/kozistr/pytorch_optimizer
 Description-Content-Type: text/markdown
 
@@ -51,7 +51,7 @@ Description-Content-Type: text/markdown
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **86 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Wide range of supported optimizers. Currently, **89 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
 * Including many variants such as `Cautious`, `AdamD`, `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
@@ -236,6 +236,9 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | APOLLO        | *SGD-like Memory, AdamW-level Performance*                                                        | [github](https://github.com/zhuhanqing/APOLLO)                                                                 | <https://arxiv.org/abs/2412.05270>                                                          | [cite](https://github.com/zhuhanqing/APOLLO?tab=readme-ov-file#-citation)                                                           |
 | MARS          | *Unleashing the Power of Variance Reduction for Training Large Models*                            | [github](https://github.com/AGI-Arena/MARS)                                                                    | <https://arxiv.org/abs/2411.10438>                                                          | [cite](https://github.com/AGI-Arena/MARS/tree/main?tab=readme-ov-file#citation)                                                     |
 | SGDSaI        | *No More Adam: Learning Rate Scaling at Initialization is All You Need*                           | [github](https://github.com/AnonymousAlethiometer/SGD_SaI)                                                     | <https://arxiv.org/abs/2411.10438>                                                          | [cite](https://github.com/AnonymousAlethiometer/SGD_SaI?tab=readme-ov-file#citation)                                                |
+| Grams         | *Gradient Descent with Adaptive Momentum Scaling*                                                 |                                                                                                                | <https://arxiv.org/abs/2412.17107>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv241217107C/exportcitation)                                                        |
+| OrthoGrad     | *Grokking at the Edge of Numerical Stability*                                                     | [github](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)                         | <https://arxiv.org/abs/2501.04697>                                                          | [cite](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability?tab=readme-ov-file#citation)                    |
+| Adam-ATAN2    | *Scaling Exponents Across Parameterizations and Optimizers*                                       |                                                                                                                | <https://arxiv.org/abs/2407.05872>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240705872E/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 
@@ -411,6 +414,10 @@ Correcting the norm of a gradient in each iteration based on the adaptive traini
 
 Updates only occur when the proposed update direction aligns with the current gradient.
 
+### Adam-ATAN2
+
+Adam-atan2 is a new numerically stable, scale-invariant version of Adam that eliminates the epsilon hyperparameter.
+
 ## Frequently asked questions
 
 [here](docs/qa.md)

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/README.md

@@ -10,7 +10,7 @@
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **86 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Wide range of supported optimizers. Currently, **89 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
 * Including many variants such as `Cautious`, `AdamD`, `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
@@ -195,6 +195,9 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | APOLLO        | *SGD-like Memory, AdamW-level Performance*                                                        | [github](https://github.com/zhuhanqing/APOLLO)                                                                 | <https://arxiv.org/abs/2412.05270>                                                          | [cite](https://github.com/zhuhanqing/APOLLO?tab=readme-ov-file#-citation)                                                           |
 | MARS          | *Unleashing the Power of Variance Reduction for Training Large Models*                            | [github](https://github.com/AGI-Arena/MARS)                                                                    | <https://arxiv.org/abs/2411.10438>                                                          | [cite](https://github.com/AGI-Arena/MARS/tree/main?tab=readme-ov-file#citation)                                                     |
 | SGDSaI        | *No More Adam: Learning Rate Scaling at Initialization is All You Need*                           | [github](https://github.com/AnonymousAlethiometer/SGD_SaI)                                                     | <https://arxiv.org/abs/2411.10438>                                                          | [cite](https://github.com/AnonymousAlethiometer/SGD_SaI?tab=readme-ov-file#citation)                                                |
+| Grams         | *Gradient Descent with Adaptive Momentum Scaling*                                                 |                                                                                                                | <https://arxiv.org/abs/2412.17107>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv241217107C/exportcitation)                                                        |
+| OrthoGrad     | *Grokking at the Edge of Numerical Stability*                                                     | [github](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)                         | <https://arxiv.org/abs/2501.04697>                                                          | [cite](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability?tab=readme-ov-file#citation)                    |
+| Adam-ATAN2    | *Scaling Exponents Across Parameterizations and Optimizers*                                       |                                                                                                                | <https://arxiv.org/abs/2407.05872>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240705872E/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 
@@ -370,6 +373,10 @@ Correcting the norm of a gradient in each iteration based on the adaptive traini
 
 Updates only occur when the proposed update direction aligns with the current gradient.
 
+### Adam-ATAN2
+
+Adam-atan2 is a new numerically stable, scale-invariant version of Adam that eliminates the epsilon hyperparameter.
+
 ## Frequently asked questions
 
 [here](docs/qa.md)

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.3.2"
+version = "3.3.3"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -14,13 +14,13 @@ keywords = [
     "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", "FAdam", "Fromage", "FTRL", "GaLore", "Gravity", "GrokFast", "GSAM", "Kate",
-    "Lamb", "LaProp", "LARS", "Lion", "LOMO", "Lookahead", "MADGRAD", "MARS", "MSVAG", "Muno", "Nero", "NovoGrad",
-    "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM",
-    "ScheduleFreeSGD", "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SGDP", "Shampoo", "ScalableShampoo", "SGDW",
-    "SignSGD", "SM3", "SOAP", "SopihaH", "SRMM", "StableAdamW", "SWATS", "Tiger", "TRAC", "WSAM", "Yogi", "BCE",
-    "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky",
-    "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
+    "DAdaptLion", "DeMo", "DiffGrad", "FAdam", "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", "QHAdam", "QHM", "RAdam", "Ranger",
+    "Ranger21", "RotoGrad", "SAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SGDP", "Shampoo",
+    "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH", "SRMM", "StableAdamW", "SWATS", "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",
@@ -97,8 +97,8 @@ select = [
     "TID", "ARG", "ERA", "RUF", "YTT", "PL", "Q"
 ]
 ignore = [
-    "B905", "D100", "D102", "D104", "D105", "D107", "D203", "D213", "D413", "PIE790", "PLR0912", "PLR0913", "PLR0915",
-    "PLR2004", "RUF013", "Q003", "ARG002",
+    "A005", "B905", "D100", "D102", "D104", "D105", "D107", "D203", "D213", "D413", "PIE790", "PLR0912", "PLR0913",
+    "PLR0915", "PLR2004", "RUF013", "Q003", "ARG002",
 ]
 fixable = ["ALL"]
 unfixable = ["F401"]

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/pytorch_optimizer/__init__.py

@@ -105,6 +105,7 @@ from pytorch_optimizer.optimizer import (
     FAdam,
     Fromage,
     GaLore,
+    Grams,
     Gravity,
     GrokFastAdamW,
     Kate,
@@ -115,6 +116,7 @@ from pytorch_optimizer.optimizer import (
     Muon,
     Nero,
     NovoGrad,
+    OrthoGrad,
     PAdam,
     PCGrad,
     Prodigy,
@@ -122,6 +124,7 @@ from pytorch_optimizer.optimizer import (
     RAdam,
     Ranger,
     Ranger21,
+    Ranger25,
     RotoGrad,
     SafeFP16Optimizer,
     ScalableShampoo,

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/pytorch_optimizer/base/optimizer.py

@@ -266,6 +266,16 @@ class BaseOptimizer(ABC, Optimizer):
         mask.mul_(mask.numel() / (mask.sum() + 1))
         update.mul_(mask)
 
+    @staticmethod
+    def get_stable_adamw_rms(grad: torch.Tensor, exp_avg_sq: torch.Tensor, eps: float = 1e-16) -> float:
+        r"""Get StableAdamW RMS.
+
+        :param grad: torch.Tensor. gradient.
+        :param exp_avg_sq: torch.Tensor. exp_avg_sq.
+        :param eps: float. epsilon.
+        """
+        return grad.pow(2).div_(exp_avg_sq.clip(min=eps)).mean().sqrt_().clip_(min=1.0).item()
+
     @staticmethod
     def validate_range(x: float, name: str, low: float, high: float, range_type: str = '[)') -> None:
         if range_type == '[)' and not low <= x < high:

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/pytorch_optimizer/optimizer/__init__.py

@@ -40,12 +40,14 @@ from pytorch_optimizer.optimizer.came import CAME
 from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptLion, DAdaptSGD
 from pytorch_optimizer.optimizer.demo import DeMo
 from pytorch_optimizer.optimizer.diffgrad import DiffGrad
+from pytorch_optimizer.optimizer.experimental.ranger25 import Ranger25
 from pytorch_optimizer.optimizer.fadam import FAdam
 from pytorch_optimizer.optimizer.fp16 import DynamicLossScaler, SafeFP16Optimizer
 from pytorch_optimizer.optimizer.fromage import Fromage
 from pytorch_optimizer.optimizer.ftrl import FTRL
 from pytorch_optimizer.optimizer.galore import GaLore
 from pytorch_optimizer.optimizer.gc import centralize_gradient
+from pytorch_optimizer.optimizer.grams import Grams
 from pytorch_optimizer.optimizer.gravity import Gravity
 from pytorch_optimizer.optimizer.grokfast import GrokFastAdamW
 from pytorch_optimizer.optimizer.kate import Kate
@@ -61,6 +63,7 @@ from pytorch_optimizer.optimizer.msvag import MSVAG
 from pytorch_optimizer.optimizer.muon import Muon
 from pytorch_optimizer.optimizer.nero import Nero
 from pytorch_optimizer.optimizer.novograd import NovoGrad
+from pytorch_optimizer.optimizer.orthograd import OrthoGrad
 from pytorch_optimizer.optimizer.padam import PAdam
 from pytorch_optimizer.optimizer.pcgrad import PCGrad
 from pytorch_optimizer.optimizer.pid import PID
@@ -282,6 +285,8 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     LaProp,
     MARS,
     SGDSaI,
+    Grams,
+    Ranger25,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/pytorch_optimizer/optimizer/adamw.py

@@ -107,13 +107,12 @@ class StableAdamW(BaseOptimizer):
                 exp_avg.lerp_(grad, weight=beta1_comp)
                 exp_avg_sq.mul_(beta2_hat).addcmul_(grad, grad, value=1.0 - beta2_hat)
 
-                rms = grad.pow(2).div_(exp_avg_sq.clip(min=eps_p2)).mean().sqrt_()
-
-                lr = group['lr'] / rms.clip(min=1.0)
+                rms = self.get_stable_adamw_rms(grad, exp_avg_sq, eps=eps_p2)
+                lr = group['lr'] / rms
 
                 self.apply_weight_decay(
                     p,
-                    p.grad,
+                    grad,
                     lr=lr,
                     weight_decay=group['weight_decay'],
                     weight_decouple=group['weight_decouple'],

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/pytorch_optimizer/optimizer/ademamix.py

@@ -19,7 +19,8 @@ class AdEMAMix(BaseOptimizer):
     :param fixed_decay: bool. fix weight decay.
     :param alpha: float. usually between 4 and 10 would work well.
     :param t_alpha_beta3: Optional[float]. total number of iterations is preferred when needed.
-    :param cautious: bool. whether to use cautious feature.
+    :param cautious: bool. whether to use the Cautious variant.
+    :param stable_adamw: bool. whether to use stable AdamW variant.
     :param eps: float. term added to the denominator to improve numerical stability.
     """
 
@@ -34,6 +35,7 @@ class AdEMAMix(BaseOptimizer):
         alpha: float = 5.0,
         t_alpha_beta3: Optional[float] = None,
         cautious: bool = False,
+        stable_adamw: bool = False,
         eps: float = 1e-8,
         **kwargs,
     ):
@@ -45,6 +47,7 @@ class AdEMAMix(BaseOptimizer):
         self.validate_non_negative(eps, 'eps')
 
         self.cautious = cautious
+        self.stable_adamw = stable_adamw
 
         defaults: DEFAULTS = {
             'lr': lr,
@@ -150,6 +153,9 @@ class AdEMAMix(BaseOptimizer):
                 if self.cautious:
                     self.apply_cautious(update, grad)
 
+                if self.stable_adamw:
+                    step_size /= self.get_stable_adamw_rms(grad, exp_avg_sq)
+
                 update.add_(exp_avg_slow, alpha=alpha_t).div_(de_nom)
 
                 p.add_(update, alpha=-step_size)

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/pytorch_optimizer/optimizer/adopt.py

@@ -18,6 +18,7 @@ class ADOPT(BaseOptimizer):
     :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
     :param fixed_decay: bool. fix weight decay.
     :param cautious: bool. whether to use the Cautious variant.
+    :param stable_adamw: bool. whether to use stable AdamW variant.
     :param eps: float. term added to the denominator to improve numerical stability.
     """
 
@@ -31,6 +32,7 @@ class ADOPT(BaseOptimizer):
         weight_decouple: bool = False,
         fixed_decay: bool = False,
         cautious: bool = False,
+        stable_adamw: bool = False,
         eps: float = 1e-6,
         **kwargs,
     ):
@@ -41,6 +43,7 @@ class ADOPT(BaseOptimizer):
 
         self.clip_lambda = clip_lambda
         self.cautious = cautious
+        self.stable_adamw = stable_adamw
 
         defaults: DEFAULTS = {
             'lr': lr,
@@ -80,6 +83,7 @@ class ADOPT(BaseOptimizer):
                 group['step'] = 1
 
             beta1, beta2 = group['betas']
+            lr: float = group['lr']
 
             for p in group['params']:
                 if p.grad is None:
@@ -98,7 +102,7 @@ class ADOPT(BaseOptimizer):
                 self.apply_weight_decay(
                     p=p,
                     grad=grad,
-                    lr=group['lr'],
+                    lr=lr,
                     weight_decay=group['weight_decay'],
                     weight_decouple=group['weight_decouple'],
                     fixed_decay=group['fixed_decay'],
@@ -127,6 +131,9 @@ class ADOPT(BaseOptimizer):
                 else:
                     update = exp_avg
 
-                p.add_(update, alpha=-group['lr'])
+                if self.stable_adamw:
+                    lr /= self.get_stable_adamw_rms(grad, exp_avg_sq)
+
+                p.add_(update, alpha=-lr)
 
         return loss

pytorch_optimizer-3.3.3/pytorch_optimizer/optimizer/experimental/ranger25.py

@@ -0,0 +1,184 @@
+import math
+from typing import Optional
+
+import torch
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class Ranger25(BaseOptimizer):
+    r"""Mixin' every fancy optimizer hacks.
+
+    :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 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 alpha: float. usually between 4 and 10 would work well.
+    :param t_alpha_beta3: Optional[float]. total number of iterations is preferred when needed.
+    :param n_sma_threshold: number of SMA threshold (recommended is 5).
+    :param cautious: bool. whether to use the Cautious variant.
+    :param stable_adamw: bool. whether to use stable AdamW variant.
+    :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, 0.9999),
+        weight_decay: float = 1e-3,
+        weight_decouple: bool = True,
+        fixed_decay: bool = False,
+        alpha: float = 5.0,
+        t_alpha_beta3: Optional[float] = None,
+        n_sma_threshold: int = 5,
+        cautious: bool = True,
+        stable_adamw: bool = True,
+        eps: float = 1e-8,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(alpha, 'alpha')
+        self.validate_non_negative(t_alpha_beta3, 't_alpha_beta3')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
+
+        self.n_sma_threshold = n_sma_threshold
+        self.cautious = cautious
+        self.stable_adamw = stable_adamw
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'alpha': alpha,
+            't_alpha_beta3': t_alpha_beta3,
+            'eps': eps,
+        }
+
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'Ranger25'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+
+                state['exp_avg'] = torch.zeros_like(p)
+                state['exp_avg_sq'] = torch.zeros_like(p)
+                state['exp_avg_slow'] = torch.zeros_like(p)
+
+    @staticmethod
+    def schedule_alpha(t_alpha_beta3: Optional[float], step: int, alpha: float) -> float:
+        return alpha if t_alpha_beta3 is None else min(step * alpha / t_alpha_beta3, alpha)
+
+    @staticmethod
+    def schedule_beta3(t_alpha_beta3: Optional[float], step: int, beta1: float, beta3: float) -> float:
+        if t_alpha_beta3 is None:
+            return beta3
+
+        log_beta1, log_beta3 = math.log(beta1), math.log(beta3)
+
+        return min(
+            math.exp(
+                log_beta1 * log_beta3 / ((1.0 - step / t_alpha_beta3) * log_beta3 + (step / t_alpha_beta3) * log_beta1)
+            ),
+            beta3,
+        )
+
+    @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, beta3 = group['betas']
+
+            bias_correction1: float = self.debias(beta1, group['step'])
+            bias_correction2_sq: float = math.sqrt(self.debias(beta2, group['step']))
+
+            step_size: float = group['lr'] / bias_correction1
+            clip: float = math.pow(group['step'], 0.25)
+
+            alpha_t: float = self.schedule_alpha(group['t_alpha_beta3'], group['step'], group['alpha'])
+            beta3_t: float = self.schedule_beta3(group['t_alpha_beta3'], group['step'], beta1, beta3)
+
+            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 len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+                    state['exp_avg_slow'] = torch.zeros_like(p)
+
+                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'],
+                )
+
+                exp_avg, exp_avg_sq, exp_avg_slow = state['exp_avg'], state['exp_avg_sq'], state['exp_avg_slow']
+
+                de_nom = exp_avg_sq.sqrt().clamp_(min=group['eps'])
+
+                normed_grad = grad.div(de_nom).clamp_(-clip, clip)
+
+                exp_avg.mul_(beta1).add_(normed_grad, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+                exp_avg_slow.mul_(beta3_t).add_(normed_grad, alpha=1.0 - beta3_t)
+
+                de_nom = exp_avg_sq.sqrt().div_(bias_correction2_sq).add_(group['eps'])
+
+                update = exp_avg.clone()
+                if self.cautious:
+                    self.apply_cautious(update, grad)
+
+                if self.stable_adamw:
+                    step_size /= self.get_stable_adamw_rms(grad, exp_avg_sq)
+
+                step_size, n_sma = self.get_rectify_step_size(
+                    is_rectify=True,
+                    step=group['step'],
+                    lr=step_size,
+                    beta2=beta2,
+                    n_sma_threshold=self.n_sma_threshold,
+                    degenerated_to_sgd=False,
+                )
+
+                update.add_(exp_avg_slow, alpha=alpha_t).div_(de_nom)
+
+                if n_sma >= self.n_sma_threshold:
+                    de_nom = exp_avg_sq.sqrt().add_(group['eps'])
+                    p.addcdiv_(update, de_nom, value=-step_size)
+                else:
+                    p.add_(update, alpha=-step_size)
+
+        return loss

pytorch_optimizer-3.3.3/pytorch_optimizer/optimizer/grams.py

@@ -0,0 +1,108 @@
+import math
+
+import torch
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class Grams(BaseOptimizer):
+    r"""Gradient Descent with Adaptive Momentum Scaling.
+
+    :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 weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. decoupled 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),
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        eps: float = 1e-6,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'eps': eps,
+        }
+
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'Grams'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+
+                state['exp_avg'] = torch.zeros_like(p)
+                state['exp_avg_sq'] = torch.zeros_like(p)
+
+    @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 len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                exp_avg.lerp_(grad, weight=beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+
+                update = (exp_avg / bias_correction1) / (exp_avg_sq / bias_correction2_sq).sqrt_().add_(group['eps'])
+                update.abs_().mul_(grad.sign())
+
+                self.apply_weight_decay(
+                    p,
+                    grad,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=False,
+                )
+
+                p.add_(update, alpha=-group['lr'])
+
+        return loss

pytorch_optimizer-3.3.3/pytorch_optimizer/optimizer/orthograd.py

@@ -0,0 +1,48 @@
+import torch
+
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import CLOSURE, LOSS, OPTIMIZER, PARAMETERS
+
+
+class OrthoGrad(BaseOptimizer):
+    r"""Grokking at the Edge of Numerical Stability.
+
+    A wrapper optimizer that projects gradients to be orthogonal to the current parameters before performing an update.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param optimizer: OPTIMIZER. base optimizer.
+    """
+
+    def __init__(self, params: PARAMETERS, optimizer: OPTIMIZER = torch.optim.AdamW, **kwargs):
+        self.eps: float = 1e-30
+
+        super().__init__(params, {})
+        self.base_optimizer = optimizer(self.param_groups, **kwargs)
+
+    def __str__(self) -> str:
+        return 'OrthoGrad'
+
+    @torch.no_grad()
+    def reset(self):
+        pass
+
+    @torch.no_grad()
+    def orthogonalize_gradients(self, params) -> None:
+        for p in params:
+            if p.grad is None:
+                continue
+
+            w = p.view(-1)
+            g = p.grad.view(-1)
+
+            proj = torch.dot(w, g).div_(torch.dot(w, w).add_(self.eps))
+            g_ortho = g.to(dtype=torch.float32, copy=True).sub_(w, alpha=proj)
+            g_ortho_scaled = g_ortho.mul_(g.norm(2).div_(g_ortho.norm(2).add_(self.eps)))
+
+            p.grad.copy_(g_ortho_scaled.view_as(p.grad))
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        for group in self.param_groups:
+            self.orthogonalize_gradients(group['params'])
+        return self.base_optimizer.step(closure)

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/pytorch_optimizer/optimizer/prodigy.py

@@ -16,7 +16,7 @@ class Prodigy(BaseOptimizer):
     :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
     :param lr: float. learning rate.
     :param betas: BETAS. betas.
-    :param beta3: float. coefficients for computing the Prodidy step-size using running averages. If set to None,
+    :param beta3: float. coefficients for computing the Prodigy step-size using running averages. If set to None,
         uses the value of square root of beta2.
     :param d0: float. initial D estimate for D-adaptation (default 1e-6). Rarely needs changing.
     :param d_coef: float. Coefficient in the expression for the estimate of d.
@@ -26,7 +26,8 @@ class Prodigy(BaseOptimizer):
     :param fixed_decay: bool. fix weight decay.
     :param bias_correction: bool. turn on Adam's bias correction.
     :param safeguard_warmup: bool. remove lr from the denominator of D estimate to avoid issues during warm-up stage.
-    :param eps: float. term added to the denominator to improve numerical stability.
+    :param eps: float. term added to the denominator to improve numerical stability. when eps is None, use atan2 rather
+        than epsilon and division for parameter updates.
     """
 
     def __init__(
@@ -43,7 +44,7 @@ class Prodigy(BaseOptimizer):
         fixed_decay: bool = False,
         bias_correction: bool = False,
         safeguard_warmup: bool = False,
-        eps: float = 1e-8,
+        eps: Optional[float] = 1e-8,
         **kwargs,
     ):
         self.validate_learning_rate(lr)
@@ -172,8 +173,6 @@ class Prodigy(BaseOptimizer):
 
                 exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
 
-                de_nom = exp_avg_sq.sqrt().add_(d * group['eps'])
-
                 self.apply_weight_decay(
                     p,
                     p.grad,
@@ -183,6 +182,13 @@ class Prodigy(BaseOptimizer):
                     fixed_decay=group['fixed_decay'],
                 )
 
-                p.addcdiv_(exp_avg, de_nom, value=-d_lr)
+                de_nom = exp_avg_sq.sqrt()
+
+                if group['eps'] is not None:
+                    de_nom.add_(d * group['eps'])
+                    p.addcdiv_(exp_avg, de_nom, value=-d_lr)
+                else:
+                    update = exp_avg.clone().atan2_(de_nom)
+                    p.add_(update, alpha=-d_lr)
 
         return loss

{pytorch_optimizer-3.3.2 → pytorch_optimizer-3.3.3}/pytorch_optimizer/optimizer/rotograd.py

@@ -81,7 +81,7 @@ class VanillaMTL(nn.Module):
             else:
                 preds.append(out_i)
 
-        return preds if len(extra_out) == 0 else preds, extra_out
+        return preds if len(extra_out) == 0 else (preds, extra_out)
 
     def backward(self, losses, backbone_loss=None, **kwargs):
         for loss in losses:
@@ -272,7 +272,7 @@ class RotateOnly(nn.Module):
             else:
                 preds.append(out_i)
 
-        return preds if len(extra_out) == 0 else preds, extra_out
+        return preds if len(extra_out) == 0 else (preds, extra_out)
 
     def backward(self, losses: Sequence[torch.Tensor], backbone_loss=None, **kwargs) -> None:
         r"""Compute the backward computations for the entire model.