Merge pull request #188 from kozistr/feature/lomo-optimizer

[Feature] Implement LOMO optimizer

Author: kozistr

README.rst

@@ -16,7 +16,7 @@ pytorch-optimizer
 
 | **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, 56 optimizers, 6 lr schedulers are supported!
+| Currently, 57 optimizers, 6 lr schedulers are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -218,6 +218,8 @@ You can check the supported optimizers with below code.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | 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 <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 <https://arxiv.org/abs/2306.09782>`__                       | `cite <https://github.com/OpenLMLab/LOMO#citation>`__                                                                |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 
 Supported LR Scheduler
 ----------------------

docs/changelogs/v2.11.0.md

@@ -4,6 +4,8 @@
 
 * Implement PAdam optimizer (#186)
   * [Closing the Generalization Gap of Adaptive Gradient Methods in Training Deep Neural Networks](https://arxiv.org/abs/1806.06763) 
+* Implement LOMO optimizer (#188)
+  * [Full Parameter Fine-tuning for Large Language Models with Limited Resources](https://arxiv.org/abs/2306.09782) 
 
 ### Diff
 

docs/optimizer_api.rst

@@ -504,3 +504,11 @@ PAdam
 
 .. autoclass:: pytorch_optimizer.PAdam
     :members:
+
+.. _LOMO:
+
+LOMO
+----
+
+.. autoclass:: pytorch_optimizer.LOMO
+    :members:

poetry.lock

@@ -13,9 +13,10 @@ keywords = [
     "pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "ASGD", "AccSGD", "AdaBelief", "AdaBound",
     "AdaDelta", "AdaFactor", "AdaMax", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdaHessian", "Adai", "AdamP",
     "AdamS", "Adan", "AggMo", "AliG", "Amos", "Apollo", "AvaGrad", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan",
-    "DAdaptSGD", "DiffGrad", "Fromage", "Gravity", "GSAM", "LARS", "Lamb", "Lion", "MADGRAD", "MSVAG", "Nero",
-    "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "RotoGrad",
-    "SAM", "SGDP", "SGDW", "SignSGD", "SM3", "SopihaH", "SRMM", "SWATS", "ScalableShampoo", "Shampoo", "Yogi",
+    "DAdaptSGD", "DiffGrad", "Fromage", "Gravity", "GSAM", "LARS", "Lamb", "Lion", "LOMO", "Lookahead", "MADGRAD",
+    "MSVAG", "Nero", "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger",
+    "Ranger21", "RotoGrad", "SAM", "SGDP", "SGDW", "SignSGD", "SM3", "SopihaH", "SRMM", "SWATS", "ScalableShampoo",
+    "Shampoo", "Yogi",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
@@ -55,18 +56,24 @@ isort = [
     { version = "^5.12.0", python = ">=3.8"}
 ]
 black = "^23.3.0"
-ruff = "^0.0.272"
-pytest = "^7.3.1"
+ruff = "^0.0.275"
+pytest = "^7.4.0"
 pytest-cov = "^4.1.0"
 
 [[tool.poetry.source]]
 name = "torch"
 url = "https://download.pytorch.org/whl/cpu"
-secondary = true
+priority = "explicit"
 
 [tool.ruff]
-select = ["A", "B", "C4", "D", "E", "F", "G", "I", "N", "S", "T", "ISC", "ICN", "W", "INP", "PIE", "T20", "RET", "SIM", "TID", "ARG", "ERA", "RUF", "YTT", "PL"]
-ignore = ["D100", "D102", "D104", "D105", "D107", "D203", "D213", "PIE790", "PLR0912", "PLR0913", "PLR0915", "PLR2004"]
+select = [
+    "A", "B", "C4", "D", "E", "F", "G", "I", "N", "S", "T", "ISC", "ICN", "W", "INP", "PIE", "T20", "RET", "SIM",
+    "TID", "ARG", "ERA", "RUF", "YTT", "PL",
+]
+ignore = [
+    "B905", "D100", "D102", "D104", "D105", "D107", "D203", "D213", "PIE790", "PLR0912", "PLR0913", "PLR0915",
+    "PLR2004", "RUF013",
+]
 fixable = ["A", "B", "C", "D", "E", "F"]
 unfixable = ["F401"]
 exclude = [
@@ -84,7 +91,7 @@ exclude = [
 ]
 line-length = 119
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
-target-version = "py39"
+target-version = "py311"
 
 [tool.ruff.per-file-ignores]
 "./hubconf.py" = ["D", "INP001"]

pyproject.toml

@@ -48,6 +48,7 @@
 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.lookahead import Lookahead
 from pytorch_optimizer.optimizer.madgrad import MADGRAD
 from pytorch_optimizer.optimizer.msvag import MSVAG
@@ -156,6 +157,7 @@
     SignSGD,
     Prodigy,
     PAdam,
+    LOMO,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

pytorch_optimizer/__init__.py

@@ -226,8 +226,8 @@ def validate_range(x: float, name: str, low: float, high: float, range_type: str
             raise ValueError(f'[-] {name} must be in the range ({low}, {high})')
 
     @staticmethod
-    def validate_non_negative(x: float, name: str):
-        if x < 0.0:
+    def validate_non_negative(x: Optional[float], name: str):
+        if x is not None and x < 0.0:
             raise ValueError(f'[-] {name} must be non-negative')
 
     @staticmethod
@@ -276,5 +276,5 @@ def validate_nus(self, nus: Union[float, Tuple[float, float]]):
             self.validate_range(nus[1], 'nu2', 0.0, 1.0, range_type='[]')
 
     @abstractmethod
-    def reset(self):
+    def reset(self):  # pragma: no cover
         raise NotImplementedError

pytorch_optimizer/base/optimizer.py

@@ -86,7 +86,7 @@ def step(self):
         return value
 
     @abstractmethod
-    def _step(self) -> float:
+    def _step(self) -> float:  # pragma: no cover
         raise NotImplementedError
 
     def get_lr(self) -> float:

pytorch_optimizer/base/scheduler.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Iterable, Literal, Optional, Tuple, Type, Union
+from typing import Callable, Dict, Iterable, Literal, Optional, Tuple, Type, Union
 
 import torch
 from torch.optim import Optimizer
@@ -7,9 +7,9 @@
 CLOSURE = Optional[Callable[[], float]]
 LOSS = Optional[float]
 BETAS = Union[Tuple[float, float], Tuple[float, float, float]]
-DEFAULTS = Dict[str, Any]
-PARAMETERS = Optional[Union[Iterable[Dict[str, Any]], Iterable[torch.Tensor]]]
-STATE = Dict[str, Any]
+DEFAULTS = Dict
+PARAMETERS = Optional[Union[Iterable[Dict], Iterable[torch.Tensor]]]
+STATE = Dict
 OPTIMIZER = Type[Optimizer]
 SCHEDULER = Type[_LRScheduler]
 

pytorch_optimizer/base/types.py

@@ -48,6 +48,7 @@ def __init__(
         self.last_overflow_iter: int = -1
         self.last_rescale_iter: int = -1
         self.overflows_since_rescale: int = 0
+        self.has_overflow_serial: bool = False
 
     def update_scale(self, overflow: bool):
         r"""Update the loss scale.