pytorch-optimizer 3.3.0__tar.gz → 3.3.1__tar.gz

{pytorch_optimizer-3.3.0 → pytorch_optimizer-3.3.1}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.1
 Name: pytorch_optimizer
-Version: 3.3.0
+Version: 3.3.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,AdamG,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdEMAMix,ADOPT,AdaHessian,Adai,Adalite,AdaLomo,AdamMini,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,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,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,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
 Author: kozistr
 Author-email: kozistr@gmail.com
 Maintainer: kozistr
@@ -51,11 +51,11 @@ Description-Content-Type: text/markdown
 
 ## The reasons why you use `pytorch-optimizer`.
 
-1. Wide range of supported optimizers. Currently, **83 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
-2. Including many variants such as `Cautious`, `AdamD`, `Gradient Centrailiaztion`
-3. Easy to use, clean, and tested codes
-4. Active maintenance
-5. Somewhat a bit more optimized compared to the original implementation
+* Wide range of supported optimizers. Currently, **85 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
+* Somewhat a bit more optimized compared to the original implementation
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -233,6 +233,8 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | MicroAdam     | *Accurate Adaptive Optimization with Low Space Overhead and Provable Convergence*                 | [github](https://github.com/IST-DASLab/MicroAdam)                                                              | <https://arxiv.org/abs/2405.15593>                                                          | [cite](https://github.com/IST-DASLab/MicroAdam?tab=readme-ov-file#citing)                                                           |
 | Muon          | *MomentUm Orthogonalized by Newton-schulz*                                                        | [github](https://github.com/KellerJordan/Muon)                                                                 | <https://x.com/kellerjordan0/status/1842300916864844014>                                    | [cite](https://github.com/KellerJordan/Muon)                                                                                        |
 | LaProp        | *Separating Momentum and Adaptivity in Adam*                                                      | [github](https://github.com/Z-T-WANG/LaProp-Optimizer)                                                         | <https://arxiv.org/abs/2002.04839>                                                          | [cite](https://github.com/Z-T-WANG/LaProp-Optimizer?tab=readme-ov-file#citation)                                                    |
+| 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)                                                     |
 
 ## Supported LR Scheduler
 

{pytorch_optimizer-3.3.0 → pytorch_optimizer-3.3.1}/README.md

@@ -10,11 +10,11 @@
 
 ## The reasons why you use `pytorch-optimizer`.
 
-1. Wide range of supported optimizers. Currently, **83 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
-2. Including many variants such as `Cautious`, `AdamD`, `Gradient Centrailiaztion`
-3. Easy to use, clean, and tested codes
-4. Active maintenance
-5. Somewhat a bit more optimized compared to the original implementation
+* Wide range of supported optimizers. Currently, **85 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
+* Somewhat a bit more optimized compared to the original implementation
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -192,6 +192,8 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | MicroAdam     | *Accurate Adaptive Optimization with Low Space Overhead and Provable Convergence*                 | [github](https://github.com/IST-DASLab/MicroAdam)                                                              | <https://arxiv.org/abs/2405.15593>                                                          | [cite](https://github.com/IST-DASLab/MicroAdam?tab=readme-ov-file#citing)                                                           |
 | Muon          | *MomentUm Orthogonalized by Newton-schulz*                                                        | [github](https://github.com/KellerJordan/Muon)                                                                 | <https://x.com/kellerjordan0/status/1842300916864844014>                                    | [cite](https://github.com/KellerJordan/Muon)                                                                                        |
 | LaProp        | *Separating Momentum and Adaptivity in Adam*                                                      | [github](https://github.com/Z-T-WANG/LaProp-Optimizer)                                                         | <https://arxiv.org/abs/2002.04839>                                                          | [cite](https://github.com/Z-T-WANG/LaProp-Optimizer?tab=readme-ov-file#citation)                                                    |
+| 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)                                                     |
 
 ## Supported LR Scheduler
 

{pytorch_optimizer-3.3.0 → pytorch_optimizer-3.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.3.0"
+version = "3.3.1"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -13,14 +13,14 @@ 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", "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", "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",
+    "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",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",

{pytorch_optimizer-3.3.0 → pytorch_optimizer-3.3.1}/pytorch_optimizer/__init__.py

@@ -42,6 +42,7 @@ from pytorch_optimizer.lr_scheduler import (
 )
 from pytorch_optimizer.optimizer import (
     ADOPT,
+    APOLLO,
     ASGD,
     BSAM,
     CAME,
@@ -50,6 +51,7 @@ from pytorch_optimizer.optimizer import (
     LARS,
     LOMO,
     MADGRAD,
+    MARS,
     MSVAG,
     PID,
     PNM,
@@ -90,7 +92,7 @@ from pytorch_optimizer.optimizer import (
     Aida,
     AliG,
     Amos,
-    Apollo,
+    ApolloDQN,
     AvaGrad,
     DAdaptAdaGrad,
     DAdaptAdam,

{pytorch_optimizer-3.3.0 → pytorch_optimizer-3.3.1}/pytorch_optimizer/optimizer/__init__.py

@@ -34,7 +34,7 @@ from pytorch_optimizer.optimizer.aggmo import AggMo
 from pytorch_optimizer.optimizer.aida import Aida
 from pytorch_optimizer.optimizer.alig import AliG
 from pytorch_optimizer.optimizer.amos import Amos
-from pytorch_optimizer.optimizer.apollo import Apollo
+from pytorch_optimizer.optimizer.apollo import APOLLO, ApolloDQN
 from pytorch_optimizer.optimizer.avagrad import AvaGrad
 from pytorch_optimizer.optimizer.came import CAME
 from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptLion, DAdaptSGD
@@ -56,6 +56,7 @@ from pytorch_optimizer.optimizer.lion import Lion
 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.mars import MARS
 from pytorch_optimizer.optimizer.msvag import MSVAG
 from pytorch_optimizer.optimizer.muon import Muon
 from pytorch_optimizer.optimizer.nero import Nero
@@ -228,7 +229,8 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     DAdaptAdan,
     AdamS,
     AdaFactor,
-    Apollo,
+    ApolloDQN,
+    APOLLO,
     SWATS,
     NovoGrad,
     Lion,
@@ -278,6 +280,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     Muon,
     ScheduleFreeRAdam,
     LaProp,
+    MARS,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

{pytorch_optimizer-3.3.0 → pytorch_optimizer-3.3.1}/pytorch_optimizer/optimizer/a2grad.py

@@ -1,5 +1,5 @@
 import math
-from typing import Optional
+from typing import Literal, Optional
 
 import torch
 
@@ -7,6 +7,8 @@ from pytorch_optimizer.base.exception import NoSparseGradientError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
 from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, PARAMETERS
 
+VARIANTS = Literal['uni', 'inc', 'exp']
+
 
 class A2Grad(BaseOptimizer):
     r"""Optimal Adaptive and Accelerated Stochastic Gradient Descent.
@@ -26,7 +28,7 @@ class A2Grad(BaseOptimizer):
         beta: float = 10.0,
         lips: float = 10.0,
         rho: float = 0.5,
-        variant: str = 'uni',
+        variant: VARIANTS = 'uni',
         **kwargs,
     ):
         self.validate_learning_rate(lr)

{pytorch_optimizer-3.3.0 → pytorch_optimizer-3.3.1}/pytorch_optimizer/optimizer/adamg.py

@@ -24,9 +24,9 @@ class AdamG(BaseOptimizer):
     def __init__(
         self,
         params: PARAMETERS,
-        lr: float = 1e-3,
+        lr: float = 1.0,
         betas: BETAS = (0.95, 0.999, 0.95),
-        p: float = 0.5,
+        p: float = 0.2,
         q: float = 0.24,
         weight_decay: float = 0.0,
         weight_decouple: bool = False,
@@ -88,8 +88,8 @@ class AdamG(BaseOptimizer):
 
             beta1, beta2, beta3 = group['betas']
 
-            bias_correction1: float = 1.0 - self.debias(beta1, group['step'])
-            bias_correction2: float = 1.0 - self.debias(beta2, group['step'])
+            bias_correction1: float = self.debias(beta1, group['step'])
+            bias_correction2: float = self.debias(beta2, group['step'])
             step_size: float = min(group['lr'], 1.0 / math.sqrt(group['step']))
 
             for p in group['params']:

{pytorch_optimizer-3.3.0 → pytorch_optimizer-3.3.1}/pytorch_optimizer/optimizer/adashift.py

@@ -17,6 +17,7 @@ class AdaShift(BaseOptimizer):
     :param keep_num: int. number of gradients used to compute first moment estimation.
     :param reduce_func: Optional[Callable]. function applied to squared gradients to further reduce the correlation.
         If None, no function is applied.
+    :param cautious: bool. whether to use cautious feature.
     :param eps: float. term added to the denominator to improve numerical stability.
     """
 
@@ -27,6 +28,7 @@ class AdaShift(BaseOptimizer):
         betas: BETAS = (0.9, 0.999),
         keep_num: int = 10,
         reduce_func: Optional[Callable] = torch.max,
+        cautious: bool = False,
         eps: float = 1e-10,
         **kwargs,
     ):
@@ -36,6 +38,7 @@ class AdaShift(BaseOptimizer):
         self.validate_non_negative(eps, 'eps')
 
         self.reduce_func: Callable = reduce_func if reduce_func is not None else lambda x: x
+        self.cautious = cautious
 
         defaults: DEFAULTS = {'lr': lr, 'betas': betas, 'keep_num': keep_num, 'eps': eps}
         super().__init__(params, defaults)
@@ -101,13 +104,16 @@ class AdaShift(BaseOptimizer):
                 exp_avg = state['exp_avg']
                 exp_avg.sub_(offset_grad, alpha=first_grad_weight).mul_(beta1).add_(grad, alpha=last_grad_weight)
 
-                reduced_grad_sq = self.reduce_func(offset_grad.mul_(offset_grad))
+                reduced_grad_sq = self.reduce_func(offset_grad.pow_(2))
 
                 exp_avg_sq = state['exp_avg_sq']
                 exp_avg_sq.mul_(beta2).add_(reduced_grad_sq, alpha=1.0 - beta2)
 
-                de_nom = exp_avg_sq.div(bias_correction).sqrt_().add_(group['eps'])
+                update = exp_avg.clone()
+                update.div_(exp_avg_sq.div(bias_correction).sqrt_().add_(group['eps']))
+                if self.cautious:
+                    self.apply_cautious(update, grad)
 
-                p.addcdiv_(exp_avg, de_nom, value=-group['lr'])
+                p.add_(update, alpha=-group['lr'])
 
         return loss

{pytorch_optimizer-3.3.0 → pytorch_optimizer-3.3.1}/pytorch_optimizer/optimizer/ademamix.py

@@ -146,9 +146,10 @@ class AdEMAMix(BaseOptimizer):
 
                 de_nom = exp_avg_sq.sqrt().div_(bias_correction2_sq).add_(group['eps'])
 
-                update = (exp_avg + alpha_t * exp_avg_slow).div_(de_nom)
                 if self.cautious:
-                    self.apply_cautious(update, grad)
+                    self.apply_cautious(exp_avg, grad)
+
+                update = (exp_avg + alpha_t * exp_avg_slow).div_(de_nom)
 
                 p.add_(update, alpha=-step_size)
 

pytorch_optimizer-3.3.1/pytorch_optimizer/optimizer/apollo.py

@@ -0,0 +1,304 @@
+import math
+from typing import Literal, Optional
+
+import numpy as np
+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
+from pytorch_optimizer.optimizer.galore_utils import GaLoreProjector
+
+SCALE_TYPE = Literal['channel', 'tensor']
+
+
+class ApolloDQN(BaseOptimizer):
+    r"""An Adaptive Parameter-wise Diagonal Quasi-Newton Method for Nonconvex Stochastic Optimization.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param init_lr: Optional[float]. initial learning rate (default lr / 1000).
+    :param beta: float. coefficient used for computing running averages of gradient.
+    :param rebound: str. rectified bound for diagonal hessian. (constant, belief).
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decay_type: str. type of weight decay. (l2, decoupled, stable).
+    :param warmup_steps: int. number of warmup steps.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-2,
+        init_lr: Optional[float] = 1e-5,
+        beta: float = 0.9,
+        rebound: str = 'constant',
+        weight_decay: float = 0.0,
+        weight_decay_type: str = 'l2',
+        warmup_steps: int = 500,
+        eps: float = 1e-4,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_range(beta, 'beta', 0.0, 1.0, range_type='[]')
+        self.validate_options(rebound, 'rebound', ['constant', 'belief'])
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_options(weight_decay_type, 'weight_decay_type', ['l2', 'decoupled', 'stable'])
+        self.validate_non_negative(eps, 'eps')
+
+        self.lr = lr
+        self.warmup_steps = warmup_steps
+        self.init_lr: float = init_lr if init_lr is not None else lr / 1000.0
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'init_lr': self.init_lr,
+            'beta': beta,
+            'rebound': rebound,
+            'weight_decay': weight_decay,
+            'weight_decay_type': weight_decay_type,
+            'eps': eps,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'ApolloDQN'
+
+    @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_grad'] = torch.zeros_like(p)
+                state['approx_hessian'] = torch.zeros_like(p)
+                state['update'] = 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
+
+            current_lr: float = (
+                group['lr']
+                if group['step'] >= self.warmup_steps
+                else (self.lr - group['init_lr']) * group['step'] / self.warmup_steps + group['init_lr']
+            )
+
+            weight_decay, eps = group['weight_decay'], group['eps']
+
+            bias_correction: float = self.debias(group['beta'], group['step'])
+            alpha: float = (1.0 - group['beta']) / bias_correction
+
+            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_grad'] = torch.zeros_like(p)
+                    state['approx_hessian'] = torch.zeros_like(p)
+                    state['update'] = torch.zeros_like(p)
+
+                if weight_decay > 0.0 and group['weight_decay_type'] == 'l2':
+                    grad.add_(p, alpha=weight_decay)
+
+                exp_avg_grad, b, d_p = state['exp_avg_grad'], state['approx_hessian'], state['update']
+
+                delta_grad = grad - exp_avg_grad
+                if group['rebound'] == 'belief':
+                    rebound = delta_grad.norm(p=np.inf)
+                else:
+                    rebound = 1e-2
+                    eps /= rebound
+
+                exp_avg_grad.add_(delta_grad, alpha=alpha)
+
+                de_nom = d_p.norm(p=4).add_(eps)
+                d_p.div_(de_nom)
+
+                v_sq = d_p.mul(d_p)
+                delta = delta_grad.div_(de_nom).mul_(d_p).sum().mul(-alpha) - b.mul(v_sq).sum()
+
+                b.addcmul_(v_sq, delta)
+
+                de_nom = b.abs().clamp_(min=rebound)
+                if group['rebound'] == 'belief':
+                    de_nom.add_(eps / alpha)
+
+                d_p.copy_(exp_avg_grad.div(de_nom))
+
+                if weight_decay > 0.0 and group['weight_decay_type'] != 'l2':
+                    if group['weight_decay_type'] == 'stable':
+                        weight_decay /= de_nom.mean().item()
+
+                    d_p.add_(p, alpha=weight_decay)
+
+                p.add_(d_p, alpha=-current_lr)
+
+        return loss
+
+
+class APOLLO(BaseOptimizer):
+    r"""SGD-like Memory, AdamW-level Performance.
+
+    :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 correct_bias: bool. Whether to correct bias in Adam.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-2,
+        betas: BETAS = (0.9, 0.999),
+        scale_type: SCALE_TYPE = 'tensor',
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        fixed_decay: bool = False,
+        correct_bias: 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,
+            'scale_type': scale_type,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'correct_bias': correct_bias,
+            'eps': eps,
+            **kwargs,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'APOLLO'
+
+    @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']
+
+            step_size: float = group['lr']
+            if group['correct_bias']:
+                bias_correction1: float = self.debias(beta1, group['step'])
+                bias_correction2_sq: float = math.sqrt(self.debias(beta2, group['step']))
+                step_size *= bias_correction2_sq / bias_correction1
+
+            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)
+
+                if 'rank' in group and p.dim() > 1:
+                    if 'projector' not in state:
+                        state['projector'] = GaLoreProjector(
+                            rank=group['rank'],
+                            update_proj_gap=group['update_proj_gap'],
+                            scale=group['scale'],
+                            projection_type=group['projection_type'],
+                        )
+
+                    grad = state['projector'].project(grad, group['step'], from_random_matrix=True)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+
+                de_nom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                norm_grad = exp_avg / de_nom
+                if 'rank' in group and p.dim() > 1:
+                    if group['scale_type'] == 'channel':
+                        norm_dim: int = 0 if norm_grad.shape[0] < norm_grad.shape[1] else 1
+                        scaling_factor = torch.norm(norm_grad, dim=norm_dim) / (torch.norm(grad, dim=norm_dim) + 1e-8)
+                        if norm_dim == 1:
+                            scaling_factor = scaling_factor.unsqueeze(1)
+                    else:
+                        scaling_factor = torch.norm(norm_grad) / (torch.norm(grad) + 1e-8)
+
+                    scaling_grad = grad * scaling_factor
+
+                    scaling_grad_norm = torch.norm(scaling_grad)
+                    if 'scaling_grad' in state:
+                        limiter = (
+                            max(
+                                scaling_grad_norm / (state['scaling_grad'] + 1e-8),
+                                1.01,
+                            )
+                            / 1.01
+                        )
+
+                        scaling_grad.div_(limiter)
+                        scaling_grad_norm.div_(limiter)
+
+                    state['scaling_grad'] = scaling_grad_norm
+
+                    norm_grad = scaling_grad * np.sqrt(group['scale'])
+                    norm_grad = state['projector'].project_back(norm_grad)
+
+                p.add_(norm_grad, alpha=-step_size)
+
+                self.apply_weight_decay(
+                    p,
+                    grad,
+                    lr=step_size,
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=group['fixed_decay'],
+                )
+
+        return loss