pytorch-optimizer 3.0.2__tar.gz → 3.1.1__tar.gz

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.1
 Name: pytorch_optimizer
-Version: 3.0.2
+Version: 3.1.1
 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,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdaHessian,Adai,Adalite,AdamMini,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DiffGrad,FAdam,Fromage,GaLore,Gravity,GrokFast,GSAM,Kate,Lamb,LARS,Lion,LOMO,Lookahead,MADGRAD,MSVAG,Nero,NovoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SopihaH,SRMM,StableAdamW,SWATS,Tiger,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD
+Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdamG,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdaHessian,Adai,Adalite,AdaLomo,AdamMini,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DiffGrad,FAdam,Fromage,GaLore,Gravity,GrokFast,GSAM,Kate,Lamb,LARS,Lion,LOMO,Lookahead,MADGRAD,MSVAG,Nero,NovoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,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
@@ -49,7 +49,7 @@ Description-Content-Type: text/markdown
 
 **pytorch-optimizer** is optimizer & lr scheduler collections in PyTorch. 
 I just re-implemented (speed & memory tweaks, plug-ins) the algorithm while based on the original paper. Also, It includes useful and practical optimization ideas.  
-Currently, **72 optimizers (+ `bitsandbytes`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
+Currently, **75 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -66,8 +66,9 @@ So, please double-check the license before using it at your work.
 $ pip3 install pytorch-optimizer
 ```
 
-From `v2.12.0`, you can install and import `bitsandbytes` optimizers. 
-please check [the requirements](https://github.com/TimDettmers/bitsandbytes?tab=readme-ov-file#tldr) before installing it.
+From `v2.12.0`, `v3.1.0`, you can use `bitsandbytes`, `q-galore-torch` optimizers respectively!
+please check [the bnb requirements](https://github.com/TimDettmers/bitsandbytes?tab=readme-ov-file#tldr), [q-galore-torch installation](https://github.com/VITA-Group/Q-GaLore?tab=readme-ov-file#install-q-galore-optimizer)
+ before installing it.
 
 From `v3.0.0`, drop `Python 3.7` support. However, you can still use this package with `Python 3.7` by installing with `--ignore-requires-python` option.
 
@@ -195,6 +196,7 @@ supported_optimizers = get_supported_optimizers()
 | Prodigy       | *An Expeditiously Adaptive Parameter-Free Learner*                                                | [github](https://github.com/konstmish/prodigy)                                                                 | <https://arxiv.org/abs/2306.06101>                                                         | [cite](https://github.com/konstmish/prodigy#how-to-cite)                                                                            |
 | PAdam         | *Closing the Generalization Gap of Adaptive Gradient Methods in Training Deep Neural Networks*    | [github](https://github.com/uclaml/Padam)                                                                      | <https://arxiv.org/abs/1806.06763>                                                         | [cite](https://github.com/uclaml/Padam#citation)                                                                                    |
 | LOMO          | *Full Parameter Fine-tuning for Large Language Models with Limited Resources*                     | [github](https://github.com/OpenLMLab/LOMO)                                                                    | <https://arxiv.org/abs/2306.09782>                                                         | [cite](https://github.com/OpenLMLab/LOMO#citation)                                                                                  |
+| AdaLOMO       | *Low-memory Optimization with Adaptive Learning Rate*                                             | [github](https://github.com/OpenLMLab/LOMO)                                                                    | <https://arxiv.org/abs/2310.10195>                                                         | [cite](https://github.com/OpenLMLab/LOMO#citation)                                                                                  |
 | Tiger         | *A Tight-fisted Optimizer, an optimizer that is extremely budget-conscious*                       | [github](https://github.com/bojone/tiger)                                                                      |                                                                                            | [cite](https://github.com/bojone/tiger/blob/main/README_en.md#citation)                                                             |
 | CAME          | *Confidence-guided Adaptive Memory Efficient Optimization*                                        | [github](https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/CAME)                            | <https://aclanthology.org/2023.acl-long.243/>                                              | [cite](https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/CAME#citation)                                          |
 | WSAM          | *Sharpness-Aware Minimization Revisited: Weighted Sharpness as a Regularization Term*             | [github](https://github.com/intelligent-machine-learning/dlrover/blob/master/atorch/atorch/optimizers/wsam.py) | <https://arxiv.org/abs/2305.15817>                                                         | [cite](https://github.com/intelligent-machine-learning/dlrover)                                                                     |
@@ -208,6 +210,8 @@ supported_optimizers = get_supported_optimizers()
 | Kate          | *Remove that Square Root: A New Efficient Scale-Invariant Version of AdaGrad*                     | [github](https://github.com/nazya/KATE)                                                                        | <https://arxiv.org/abs/2403.02648>                                                         | [cite](https://github.com/nazya/KATE?tab=readme-ov-file#remove-that-square-root-a-new-efficient-scale-invariant-version-of-adagrad) |
 | StableAdamW   | *Stable and low-precision training for large-scale vision-language models*                        |                                                                                                                | <https://arxiv.org/abs/2304.13013>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2023arXiv230413013W/exportcitation)                                                        |
 | AdamMini      | *Use Fewer Learning Rates To Gain More*                                                           | [github](https://github.com/zyushun/Adam-mini)                                                                 | <https://arxiv.org/abs/2406.16793>                                                         | [cite](https://github.com/zyushun/Adam-mini?tab=readme-ov-file#citation)                                                            |
+| TRAC          | *Adaptive Parameter-free Optimization*                                                            | [github](https://github.com/ComputationalRobotics/TRAC)                                                        | <https://arxiv.org/abs/2405.16642>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240516642M/exportcitation)                                                        |
+| AdamG         | *Towards Stability of Parameter-free Optimization*                                                |                                                                                                                | <https://arxiv.org/abs/2405.04376>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240504376P/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 
@@ -357,7 +361,11 @@ Correcting the norm of a gradient in each iteration based on the adaptive traini
 
 ## Frequently asked questions
 
-[here](./qa.md)
+[here](docs/qa.md)
+
+## Visualization
+
+[here](docs/visualization.md)
 
 ## Citation
 
@@ -369,7 +377,7 @@ If you use this software, please cite it below. Or you can get it from "cite thi
         month = jan,
         title = {{pytorch_optimizer: optimizer & lr scheduler & loss function collections in PyTorch}},
         url = {https://github.com/kozistr/pytorch_optimizer},
-        version = {3.0.1},
+        version = {3.1.0},
         year = {2021}
     }
 

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/README.md

@@ -10,7 +10,7 @@
 
 **pytorch-optimizer** is optimizer & lr scheduler collections in PyTorch. 
 I just re-implemented (speed & memory tweaks, plug-ins) the algorithm while based on the original paper. Also, It includes useful and practical optimization ideas.  
-Currently, **72 optimizers (+ `bitsandbytes`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
+Currently, **75 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -27,8 +27,9 @@ So, please double-check the license before using it at your work.
 $ pip3 install pytorch-optimizer
 ```
 
-From `v2.12.0`, you can install and import `bitsandbytes` optimizers. 
-please check [the requirements](https://github.com/TimDettmers/bitsandbytes?tab=readme-ov-file#tldr) before installing it.
+From `v2.12.0`, `v3.1.0`, you can use `bitsandbytes`, `q-galore-torch` optimizers respectively!
+please check [the bnb requirements](https://github.com/TimDettmers/bitsandbytes?tab=readme-ov-file#tldr), [q-galore-torch installation](https://github.com/VITA-Group/Q-GaLore?tab=readme-ov-file#install-q-galore-optimizer)
+ before installing it.
 
 From `v3.0.0`, drop `Python 3.7` support. However, you can still use this package with `Python 3.7` by installing with `--ignore-requires-python` option.
 
@@ -156,6 +157,7 @@ supported_optimizers = get_supported_optimizers()
 | Prodigy       | *An Expeditiously Adaptive Parameter-Free Learner*                                                | [github](https://github.com/konstmish/prodigy)                                                                 | <https://arxiv.org/abs/2306.06101>                                                         | [cite](https://github.com/konstmish/prodigy#how-to-cite)                                                                            |
 | PAdam         | *Closing the Generalization Gap of Adaptive Gradient Methods in Training Deep Neural Networks*    | [github](https://github.com/uclaml/Padam)                                                                      | <https://arxiv.org/abs/1806.06763>                                                         | [cite](https://github.com/uclaml/Padam#citation)                                                                                    |
 | LOMO          | *Full Parameter Fine-tuning for Large Language Models with Limited Resources*                     | [github](https://github.com/OpenLMLab/LOMO)                                                                    | <https://arxiv.org/abs/2306.09782>                                                         | [cite](https://github.com/OpenLMLab/LOMO#citation)                                                                                  |
+| AdaLOMO       | *Low-memory Optimization with Adaptive Learning Rate*                                             | [github](https://github.com/OpenLMLab/LOMO)                                                                    | <https://arxiv.org/abs/2310.10195>                                                         | [cite](https://github.com/OpenLMLab/LOMO#citation)                                                                                  |
 | Tiger         | *A Tight-fisted Optimizer, an optimizer that is extremely budget-conscious*                       | [github](https://github.com/bojone/tiger)                                                                      |                                                                                            | [cite](https://github.com/bojone/tiger/blob/main/README_en.md#citation)                                                             |
 | CAME          | *Confidence-guided Adaptive Memory Efficient Optimization*                                        | [github](https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/CAME)                            | <https://aclanthology.org/2023.acl-long.243/>                                              | [cite](https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/CAME#citation)                                          |
 | WSAM          | *Sharpness-Aware Minimization Revisited: Weighted Sharpness as a Regularization Term*             | [github](https://github.com/intelligent-machine-learning/dlrover/blob/master/atorch/atorch/optimizers/wsam.py) | <https://arxiv.org/abs/2305.15817>                                                         | [cite](https://github.com/intelligent-machine-learning/dlrover)                                                                     |
@@ -169,6 +171,8 @@ supported_optimizers = get_supported_optimizers()
 | Kate          | *Remove that Square Root: A New Efficient Scale-Invariant Version of AdaGrad*                     | [github](https://github.com/nazya/KATE)                                                                        | <https://arxiv.org/abs/2403.02648>                                                         | [cite](https://github.com/nazya/KATE?tab=readme-ov-file#remove-that-square-root-a-new-efficient-scale-invariant-version-of-adagrad) |
 | StableAdamW   | *Stable and low-precision training for large-scale vision-language models*                        |                                                                                                                | <https://arxiv.org/abs/2304.13013>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2023arXiv230413013W/exportcitation)                                                        |
 | AdamMini      | *Use Fewer Learning Rates To Gain More*                                                           | [github](https://github.com/zyushun/Adam-mini)                                                                 | <https://arxiv.org/abs/2406.16793>                                                         | [cite](https://github.com/zyushun/Adam-mini?tab=readme-ov-file#citation)                                                            |
+| TRAC          | *Adaptive Parameter-free Optimization*                                                            | [github](https://github.com/ComputationalRobotics/TRAC)                                                        | <https://arxiv.org/abs/2405.16642>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240516642M/exportcitation)                                                        |
+| AdamG         | *Towards Stability of Parameter-free Optimization*                                                |                                                                                                                | <https://arxiv.org/abs/2405.04376>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240504376P/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 
@@ -318,7 +322,11 @@ Correcting the norm of a gradient in each iteration based on the adaptive traini
 
 ## Frequently asked questions
 
-[here](./qa.md)
+[here](docs/qa.md)
+
+## Visualization
+
+[here](docs/visualization.md)
 
 ## Citation
 
@@ -330,7 +338,7 @@ If you use this software, please cite it below. Or you can get it from "cite thi
         month = jan,
         title = {{pytorch_optimizer: optimizer & lr scheduler & loss function collections in PyTorch}},
         url = {https://github.com/kozistr/pytorch_optimizer},
-        version = {3.0.1},
+        version = {3.1.0},
         year = {2021}
     }
 

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.0.2"
+version = "3.1.1"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -11,14 +11,15 @@ 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", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdaHessian", "Adai", "Adalite",
-    "AdamMini", "AdamP", "AdamS", "Adan", "AggMo", "Aida", "AliG", "Amos", "Apollo", "AvaGrad", "bSAM", "CAME",
-    "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DiffGrad", "FAdam", "Fromage", "GaLore",
-    "Gravity", "GrokFast", "GSAM", "Kate", "Lamb", "LARS", "Lion", "LOMO", "Lookahead", "MADGRAD", "MSVAG", "Nero",
-    "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "RotoGrad",
-    "SAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3",
-    "SopihaH", "SRMM", "StableAdamW", "SWATS", "Tiger", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine",
-    "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD",
+    "AdaDelta", "AdaFactor", "AdaMax", "AdamG", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdaHessian", "Adai",
+    "Adalite", "AdaLomo", "AdamMini", "AdamP", "AdamS", "Adan", "AggMo", "Aida", "AliG", "Amos", "Apollo", "AvaGrad",
+    "bSAM", "CAME", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DiffGrad", "FAdam",
+    "Fromage", "GaLore", "Gravity", "GrokFast", "GSAM", "Kate", "Lamb", "LARS", "Lion", "LOMO", "Lookahead", "MADGRAD",
+    "MSVAG", "Nero", "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger",
+    "Ranger21", "RotoGrad", "SAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "SGDP", "Shampoo", "ScalableShampoo", "SGDW",
+    "SignSGD", "SM3", "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",
@@ -74,6 +75,10 @@ lint.ignore = [
 ]
 lint.fixable = ["ALL"]
 lint.unfixable = ["F401"]
+lint.dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+lint.flake8-quotes.docstring-quotes = "double"
+lint.flake8-quotes.inline-quotes = "single"
+lint.pylint.max-args = 7
 exclude = [
     ".git",
     ".github",
@@ -88,11 +93,7 @@ exclude = [
     "__pypackages__",
 ]
 line-length = 119
-lint.dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 target-version = "py311"
-lint.flake8-quotes.docstring-quotes = "double"
-lint.flake8-quotes.inline-quotes = "single"
-lint.pylint.max-args = 7
 
 [tool.ruff.format]
 quote-style = "single"
@@ -101,6 +102,7 @@ quote-style = "single"
 "./pytorch_optimizer/__init__.py" = ["F401"]
 "./pytorch_optimizer/lr_scheduler/__init__.py" = ["F401"]
 "./hubconf.py" = ["D", "INP001"]
+"./examples/visualize_optimizers.py" = ["D103", "D400", "D415"]
 "./tests/__init__.py" = ["D"]
 "./tests/constants.py" = ["D"]
 "./tests/utils.py" = ["D"]

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/pytorch_optimizer/__init__.py

@@ -1,8 +1,10 @@
 # ruff: noqa
+from importlib.util import find_spec
 from typing import Dict, List
 
 import torch.cuda
 from torch import nn
+from torch.optim import AdamW
 
 from pytorch_optimizer.base.types import OPTIMIZER, PARAMETERS, SCHEDULER
 from pytorch_optimizer.loss.bi_tempered import BinaryBiTemperedLogisticLoss, BiTemperedLogisticLoss
@@ -42,6 +44,7 @@ from pytorch_optimizer.optimizer.adai import Adai
 from pytorch_optimizer.optimizer.adalite import Adalite
 from pytorch_optimizer.optimizer.adam_mini import AdamMini
 from pytorch_optimizer.optimizer.adamax import AdaMax
+from pytorch_optimizer.optimizer.adamg import AdamG
 from pytorch_optimizer.optimizer.adamod import AdaMod
 from pytorch_optimizer.optimizer.adamp import AdamP
 from pytorch_optimizer.optimizer.adams import AdamS
@@ -72,7 +75,7 @@ from pytorch_optimizer.optimizer.kate import Kate
 from pytorch_optimizer.optimizer.lamb import Lamb
 from pytorch_optimizer.optimizer.lars import LARS
 from pytorch_optimizer.optimizer.lion import Lion
-from pytorch_optimizer.optimizer.lomo import LOMO
+from pytorch_optimizer.optimizer.lomo import LOMO, AdaLOMO
 from pytorch_optimizer.optimizer.lookahead import Lookahead
 from pytorch_optimizer.optimizer.madgrad import MADGRAD
 from pytorch_optimizer.optimizer.msvag import MSVAG
@@ -114,6 +117,7 @@ from pytorch_optimizer.optimizer.sophia import SophiaH
 from pytorch_optimizer.optimizer.srmm import SRMM
 from pytorch_optimizer.optimizer.swats import SWATS
 from pytorch_optimizer.optimizer.tiger import Tiger
+from pytorch_optimizer.optimizer.trac import TRAC
 from pytorch_optimizer.optimizer.utils import (
     clip_grad_norm,
     disable_running_stats,
@@ -126,14 +130,11 @@ from pytorch_optimizer.optimizer.utils import (
 )
 from pytorch_optimizer.optimizer.yogi import Yogi
 
-try:
-    import bitsandbytes as bnb
-
-    HAS_BNB: bool = True  # pragma: no cover
-except ImportError:
-    HAS_BNB: bool = False
+HAS_BNB: bool = find_spec('bitsandbytes') is not None
+HAS_Q_GALORE: bool = find_spec('q-galore-torch') is not None
 
 OPTIMIZER_LIST: List[OPTIMIZER] = [
+    AdamW,
     AdaBelief,
     AdaBound,
     PID,
@@ -205,6 +206,8 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     Kate,
     StableAdamW,
     AdamMini,
+    AdaLOMO,
+    AdamG,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 
@@ -252,22 +255,58 @@ LOSS_FUNCTIONS: Dict[str, nn.Module] = {
 
 def load_bnb_optimizer(optimizer: str) -> OPTIMIZER:  # pragma: no cover
     r"""load bnb optimizer instance."""
+    from bitsandbytes import optim
+
     if 'sgd8bit' in optimizer:
-        return bnb.optim.SGD8bit
+        return optim.SGD8bit
     if 'adam8bit' in optimizer:
-        return bnb.optim.Adam8bit
+        return optim.Adam8bit
+    if 'paged_adam8bit' in optimizer:
+        return optim.PagedAdam8bit
     if 'adamw8bit' in optimizer:
-        return bnb.optim.AdamW8bit
+        return optim.AdamW8bit
+    if 'paged_adamw8bit' in optimizer:
+        return optim.PagedAdamW8bit
     if 'lamb8bit' in optimizer:
-        return bnb.optim.LAMB8bit
+        return optim.LAMB8bit
     if 'lars8bit' in optimizer:
-        return bnb.optim.LARS8bit
+        return optim.LARS8bit
     if 'lion8bit' in optimizer:
-        return bnb.optim.Lion8bit
+        return optim.Lion8bit
     if 'adagrad8bit' in optimizer:
-        return bnb.optim.Adagrad8bit
+        return optim.Adagrad8bit
     if 'rmsprop8bit' in optimizer:
-        return bnb.optim.RMSprop8bit
+        return optim.RMSprop8bit
+    if 'adagrad32bit' in optimizer:
+        return optim.Adagrad32bit
+    if 'adam32bit' in optimizer:
+        return optim.Adam32bit
+    if 'paged_adam32bit' in optimizer:
+        return optim.PagedAdam32bit
+    if 'adamw32bit' in optimizer:
+        return optim.AdamW32bit
+    if 'lamb32bit' in optimizer:
+        return optim.LAMB32bit
+    if 'lars32bit' in optimizer:
+        return optim.LARS32bit
+    if 'lion32bit' in optimizer:
+        return optim.Lion32bit
+    if 'paged_lion32bit' in optimizer:
+        return optim.PagedLion32bit
+    if 'rmsprop32bit' in optimizer:
+        return optim.RMSprop32bit
+    if 'sgd32bit' in optimizer:
+        return optim.SGD32bit
+    raise NotImplementedError(f'[-] not implemented optimizer : {optimizer}')
+
+
+def load_q_galore_optimizer(optimizer: str) -> OPTIMIZER:  # pragma: no cover
+    r"""load Q-GaLore optimizer instance."""
+    import q_galore_torch
+
+    if 'adamw8bit' in optimizer:
+        return q_galore_torch.QGaLoreAdamW8bit
+
     raise NotImplementedError(f'[-] not implemented optimizer : {optimizer}')
 
 
@@ -277,7 +316,11 @@ def load_optimizer(optimizer: str) -> OPTIMIZER:
     if optimizer.startswith('bnb'):
         if HAS_BNB and torch.cuda.is_available():
             return load_bnb_optimizer(optimizer)  # pragma: no cover
-        raise ImportError(f'[-] bitsandbytes and CUDA required for bnb optimizers : {optimizer}')
+        raise ImportError(f'[-] bitsandbytes and CUDA required for the optimizer {optimizer}')
+    if optimizer.startswith('q_galore'):
+        if HAS_Q_GALORE and torch.cuda.is_available():
+            return load_q_galore_optimizer(optimizer)  # pragma: no cover
+        raise ImportError(f'[-] bitsandbytes, q-galore-torch, and CUDA required for the optimizer {optimizer}')
     if optimizer not in OPTIMIZERS:
         raise NotImplementedError(f'[-] not implemented optimizer : {optimizer}')
 
@@ -312,6 +355,8 @@ def create_optimizer(
 
     if optimizer_name == 'alig':
         optimizer = optimizer(parameters, max_lr=lr, **kwargs)
+    elif optimizer_name in {'lomo', 'adalomo', 'adammini'}:
+        optimizer = optimizer(model, lr=lr, **kwargs)
     else:
         optimizer = optimizer(parameters, lr=lr, **kwargs)
 

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/pytorch_optimizer/base/optimizer.py

@@ -3,17 +3,21 @@ from abc import ABC, abstractmethod
 from typing import List, Optional, Tuple, Union
 
 import torch
+from torch.optim import Optimizer
 
 from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError
-from pytorch_optimizer.base.types import BETAS, HUTCHINSON_G, PARAMETERS, STATE
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, HUTCHINSON_G, LOSS, PARAMETERS, STATE
 
 
-class BaseOptimizer(ABC):
-    r"""Base optimizer class."""
+class BaseOptimizer(ABC, Optimizer):
+    r"""Base optimizer class. Provides common functionalities for the optimizers."""
+
+    def __init__(self, params: PARAMETERS, defaults: DEFAULTS) -> None:
+        super().__init__(params, defaults)
 
     @staticmethod
     @torch.no_grad()
-    def set_hessian(param_groups: PARAMETERS, state: STATE, hessian: List[torch.Tensor]):
+    def set_hessian(param_groups: PARAMETERS, state: STATE, hessian: List[torch.Tensor]) -> None:
         r"""Set hessian to state from external source. Generally useful when using functorch as a base.
 
         Example:
@@ -45,7 +49,7 @@ class BaseOptimizer(ABC):
                 i += 1
 
     @staticmethod
-    def zero_hessian(param_groups: PARAMETERS, state: STATE, pre_zero: bool = True):
+    def zero_hessian(param_groups: PARAMETERS, state: STATE, pre_zero: bool = True) -> None:
         r"""Zero-out hessian.
 
         :param param_groups: PARAMETERS. parameter groups.
@@ -68,7 +72,7 @@ class BaseOptimizer(ABC):
         num_samples: int = 1,
         alpha: float = 1.0,
         distribution: HUTCHINSON_G = 'gaussian',
-    ):
+    ) -> None:
         r"""Hutchinson's approximate hessian, added to the state under key `hessian`.
 
         :param param_groups: PARAMETERS. parameter groups.
@@ -110,7 +114,7 @@ class BaseOptimizer(ABC):
         weight_decouple: bool,
         fixed_decay: bool,
         ratio: Optional[float] = None,
-    ):
+    ) -> None:
         r"""Apply weight decay.
 
         :param p: torch.Tensor. parameter.
@@ -145,6 +149,27 @@ class BaseOptimizer(ABC):
 
         return de_nom.sqrt_().add_(eps)
 
+    @staticmethod
+    def debias(beta: float, step: int) -> float:
+        r"""Adam-style debias correction. Returns `1.0 - beta ** step`.
+
+        :param beta: float. beta.
+        :param step: int. number of step.
+        """
+        return 1.0 - math.pow(beta, step)  # fmt: skip
+
+    @staticmethod
+    def debias_beta(beta: float, step: int) -> float:
+        r"""Apply the Adam-style debias correction into beta.
+
+        Simplified version of `\^{beta} = beta * (1.0 - beta ** (step - 1)) / (1.0 - beta ** step)`
+
+        :param beta: float. beta.
+        :param step: int. number of step.
+        """
+        beta_n: float = math.pow(beta, step)
+        return (beta_n - beta) / (beta_n - 1.0)  # fmt: skip
+
     @staticmethod
     def apply_adam_debias(adam_debias: bool, step_size: float, bias_correction1: float) -> float:
         r"""Apply AdamD variant.
@@ -205,14 +230,30 @@ class BaseOptimizer(ABC):
         :param exp_grad_norm: Optional[torch.Tensor]. exp_grad_norm.
         :param r: float. Optional[float]. momentum (ratio).
         """
-        if not adanorm:
+        if not adanorm or exp_grad_norm is None:
             return grad
 
         grad_norm = torch.linalg.norm(grad)
 
         exp_grad_norm.mul_(r).add_(grad_norm, alpha=1.0 - r)
 
-        return grad * exp_grad_norm / grad_norm if exp_grad_norm > grad_norm else grad
+        return grad.mul(exp_grad_norm).div_(grad_norm) if exp_grad_norm > grad_norm else grad
+
+    @staticmethod
+    def get_rms(x: torch.Tensor) -> float:
+        r"""Get RMS."""
+        return x.norm(2) / math.sqrt(x.numel())
+
+    @staticmethod
+    def approximate_sq_grad(
+        exp_avg_sq_row: torch.Tensor,
+        exp_avg_sq_col: torch.Tensor,
+        output: torch.Tensor,
+    ) -> None:
+        r"""Get approximation of EMA of squared gradient."""
+        r_factor: torch.Tensor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
+        c_factor: torch.Tensor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+        torch.mul(r_factor, c_factor, out=output)
 
     @staticmethod
     def validate_range(x: float, name: str, low: float, high: float, range_type: str = '[)') -> None:
@@ -283,5 +324,8 @@ class BaseOptimizer(ABC):
             self.validate_range(nus[1], 'nu2', 0.0, 1.0, range_type='[]')
 
     @abstractmethod
-    def reset(self):  # pragma: no cover
+    def reset(self) -> None:  # pragma: no cover
+        raise NotImplementedError
+
+    def step(self, closure: CLOSURE = None) -> LOSS:  # pragma: no cover
         raise NotImplementedError

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/pytorch_optimizer/optimizer/a2grad.py

@@ -2,14 +2,13 @@ import math
 from typing import Optional
 
 import torch
-from torch.optim.optimizer import Optimizer
 
 from pytorch_optimizer.base.exception import NoSparseGradientError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
 from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, PARAMETERS
 
 
-class A2Grad(Optimizer, BaseOptimizer):
+class A2Grad(BaseOptimizer):
     r"""Optimal Adaptive and Accelerated Stochastic Gradient Descent.
 
     :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/pytorch_optimizer/optimizer/adabelief.py

@@ -1,14 +1,13 @@
 import math
 
 import torch
-from torch.optim.optimizer import Optimizer
 
 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 AdaBelief(Optimizer, BaseOptimizer):
+class AdaBelief(BaseOptimizer):
     r"""Adapting Step-sizes by the Belief in Observed Gradients.
 
     :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
@@ -101,8 +100,8 @@ class AdaBelief(Optimizer, BaseOptimizer):
 
             beta1, beta2 = group['betas']
 
-            bias_correction1: float = 1.0 - beta1 ** group['step']
-            bias_correction2_sq: float = math.sqrt(1.0 - beta2 ** group['step'])
+            bias_correction1: float = self.debias(beta1, group['step'])
+            bias_correction2_sq: float = math.sqrt(self.debias(beta2, group['step']))
 
             step_size, n_sma = self.get_rectify_step_size(
                 is_rectify=group['rectify'],

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/pytorch_optimizer/optimizer/adabound.py

@@ -2,14 +2,13 @@ import math
 from typing import List
 
 import torch
-from torch.optim.optimizer import Optimizer
 
 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 AdaBound(Optimizer, BaseOptimizer):
+class AdaBound(BaseOptimizer):
     r"""Adaptive Gradient Methods with Dynamic Bound of Learning Rate.
 
     :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
@@ -90,8 +89,8 @@ class AdaBound(Optimizer, BaseOptimizer):
 
             beta1, beta2 = group['betas']
 
-            bias_correction1: float = 1.0 - beta1 ** group['step']
-            bias_correction2_sq: float = math.sqrt(1.0 - beta2 ** group['step'])
+            bias_correction1: float = self.debias(beta1, group['step'])
+            bias_correction2_sq: float = math.sqrt(self.debias(beta2, group['step']))
 
             final_lr: float = group['final_lr'] * group['lr'] / base_lr
             lower_bound: float = final_lr * (1 - 1 / (group['gamma'] * group['step'] + 1))

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/pytorch_optimizer/optimizer/adadelta.py

@@ -1,12 +1,11 @@
 import torch
-from torch.optim.optimizer import Optimizer
 
 from pytorch_optimizer.base.exception import NoSparseGradientError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
 from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, PARAMETERS
 
 
-class AdaDelta(Optimizer, BaseOptimizer):
+class AdaDelta(BaseOptimizer):
     r"""An Adaptive Learning Rate Method.
 
     :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/pytorch_optimizer/optimizer/adafactor.py

@@ -2,14 +2,13 @@ import math
 from typing import Optional, Tuple
 
 import torch
-from torch.optim.optimizer import Optimizer
 
 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 AdaFactor(Optimizer, BaseOptimizer):
+class AdaFactor(BaseOptimizer):
     r"""Adaptive Learning Rates with Sublinear Memory Cost with some tweaks.
 
     :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
@@ -127,22 +126,6 @@ class AdaFactor(Optimizer, BaseOptimizer):
         r"""Get `factored`."""
         return len(shape) >= 2
 
-    @staticmethod
-    def get_rms(x: torch.Tensor) -> float:
-        r"""Get RMS."""
-        return x.norm(2) / math.sqrt(x.numel())
-
-    @staticmethod
-    def approximate_sq_grad(
-        exp_avg_sq_row: torch.Tensor,
-        exp_avg_sq_col: torch.Tensor,
-        output: torch.Tensor,
-    ):
-        r"""Get approximation of EMA of squared gradient."""
-        r_factor: torch.Tensor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
-        c_factor: torch.Tensor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
-        torch.mul(r_factor, c_factor, out=output)
-
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
         loss: LOSS = None

{pytorch_optimizer-3.0.2 → pytorch_optimizer-3.1.1}/pytorch_optimizer/optimizer/adahessian.py

@@ -1,14 +1,13 @@
 from typing import List, Optional
 
 import torch
-from torch.optim.optimizer import Optimizer
 
 from pytorch_optimizer.base.exception import NoSparseGradientError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
 from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, HUTCHINSON_G, LOSS, PARAMETERS
 
 
-class AdaHessian(Optimizer, BaseOptimizer):
+class AdaHessian(BaseOptimizer):
     r"""An Adaptive Second Order Optimizer for Machine Learning.
 
         Requires `loss.backward(create_graph=True)` in order to calculate hessians.
@@ -104,8 +103,8 @@ class AdaHessian(Optimizer, BaseOptimizer):
 
             beta1, beta2 = group['betas']
 
-            bias_correction1: float = 1.0 - beta1 ** group['step']
-            bias_correction2: float = 1.0 - beta2 ** group['step']
+            bias_correction1: float = self.debias(beta1, group['step'])
+            bias_correction2: float = self.debias(beta2, group['step'])
 
             for p in group['params']:
                 if p.grad is None: