pytorch-optimizer 2.6.0__tar.gz → 2.7.0__tar.gz

{pytorch_optimizer-2.6.0 → pytorch_optimizer-2.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytorch-optimizer
-Version: 2.6.0
+Version: 2.7.0
 Summary: optimizer & lr scheduler implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas.
 Home-page: https://github.com/kozistr/pytorch_optimizer
 License: Apache-2.0
@@ -183,6 +183,10 @@ You can check the supported optimizers & lr schedulers.
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | SM3          | *Memory-Efficient Adaptive Optimization*                                                        | `github <https://github.com/google-research/google-research/tree/master/sm3>`__   | `https://arxiv.org/abs/1901.11150 <https://arxiv.org/abs/1901.11150>`__                       |
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| AdaNorm      | *Adaptive Gradient Norm Correction based Optimizer for CNNs*                                    | `github <https://github.com/shivram1987/AdaNorm>`__                               | `https://arxiv.org/abs/2210.06364 <https://arxiv.org/abs/2210.06364>`__                       |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| RotoGrad     | *Gradient Homogenization in Multitask Learning*                                                 | `github <https://github.com/adrianjav/rotograd>`__                                | `https://openreview.net/pdf?id=T8wHz4rnuGL <https://openreview.net/pdf?id=T8wHz4rnuGL>`__     |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -390,6 +394,10 @@ Citations
 
 `SM3 <https://ui.adsabs.harvard.edu/abs/2019arXiv190111150A/exportcitation>`__
 
+`AdaNorm <https://github.com/shivram1987/AdaNorm/tree/main#citation>`__
+
+`RotoGrad <https://github.com/adrianjav/rotograd#citing>`__
+
 Citation
 --------
 

{pytorch_optimizer-2.6.0 → pytorch_optimizer-2.7.0}/README.rst

@@ -140,6 +140,10 @@ You can check the supported optimizers & lr schedulers.
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | SM3          | *Memory-Efficient Adaptive Optimization*                                                        | `github <https://github.com/google-research/google-research/tree/master/sm3>`__   | `https://arxiv.org/abs/1901.11150 <https://arxiv.org/abs/1901.11150>`__                       |
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| AdaNorm      | *Adaptive Gradient Norm Correction based Optimizer for CNNs*                                    | `github <https://github.com/shivram1987/AdaNorm>`__                               | `https://arxiv.org/abs/2210.06364 <https://arxiv.org/abs/2210.06364>`__                       |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| RotoGrad     | *Gradient Homogenization in Multitask Learning*                                                 | `github <https://github.com/adrianjav/rotograd>`__                                | `https://openreview.net/pdf?id=T8wHz4rnuGL <https://openreview.net/pdf?id=T8wHz4rnuGL>`__     |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -347,6 +351,10 @@ Citations
 
 `SM3 <https://ui.adsabs.harvard.edu/abs/2019arXiv190111150A/exportcitation>`__
 
+`AdaNorm <https://github.com/shivram1987/AdaNorm/tree/main#citation>`__
+
+`RotoGrad <https://github.com/adrianjav/rotograd#citing>`__
+
 Citation
 --------
 

{pytorch_optimizer-2.6.0 → pytorch_optimizer-2.7.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.6.0"
+version = "2.7.0"
 description = "optimizer & lr scheduler implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas."
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -103,6 +103,7 @@ target-version = "py39"
 omit = [
     "./pytorch_optimizer/optimizer/gsam.py",
     "./pytorch_optimizer/optimizer/fp16.py",
+    "./pytorch_optimizer/optimizer/rotograd.py",
 ]
 
 [build-system]

{pytorch_optimizer-2.6.0 → pytorch_optimizer-2.7.0}/pytorch_optimizer/__init__.py

@@ -23,13 +23,13 @@ from pytorch_optimizer.optimizer.adai import Adai
 from pytorch_optimizer.optimizer.adamp import AdamP
 from pytorch_optimizer.optimizer.adams import AdamS
 from pytorch_optimizer.optimizer.adan import Adan
+from pytorch_optimizer.optimizer.adanorm import AdaNorm
 from pytorch_optimizer.optimizer.adapnm import AdaPNM
 from pytorch_optimizer.optimizer.agc import agc
 from pytorch_optimizer.optimizer.alig import AliG
 from pytorch_optimizer.optimizer.apollo import Apollo
-from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptSGD
+from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptSGD
 from pytorch_optimizer.optimizer.diffgrad import DiffGrad
-from pytorch_optimizer.optimizer.diffrgrad import DiffRGrad
 from pytorch_optimizer.optimizer.fp16 import DynamicLossScaler, SafeFP16Optimizer
 from pytorch_optimizer.optimizer.gc import centralize_gradient
 from pytorch_optimizer.optimizer.gsam import GSAM
@@ -43,9 +43,9 @@ from pytorch_optimizer.optimizer.novograd import NovoGrad
 from pytorch_optimizer.optimizer.pcgrad import PCGrad
 from pytorch_optimizer.optimizer.pnm import PNM
 from pytorch_optimizer.optimizer.radam import RAdam
-from pytorch_optimizer.optimizer.ralamb import RaLamb
 from pytorch_optimizer.optimizer.ranger import Ranger
 from pytorch_optimizer.optimizer.ranger21 import Ranger21
+from pytorch_optimizer.optimizer.rotograd import RotoGrad
 from pytorch_optimizer.optimizer.sam import SAM
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo
@@ -69,8 +69,10 @@ from pytorch_optimizer.optimizer.utils import (
     clip_grad_norm,
     disable_running_stats,
     enable_running_stats,
+    get_global_gradient_norm,
     get_optimizer_parameters,
     normalize_gradient,
+    reduce_max_except_dim,
     unit_norm,
 )
 
@@ -82,14 +84,12 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     Adan,
     AdaPNM,
     DiffGrad,
-    DiffRGrad,
     Lamb,
     LARS,
     MADGRAD,
     Nero,
     PNM,
     RAdam,
-    RaLamb,
     Ranger,
     Ranger21,
     SGDP,
@@ -98,6 +98,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     DAdaptAdaGrad,
     DAdaptAdam,
     DAdaptSGD,
+    DAdaptAdan,
     AdamS,
     AdaFactor,
     Apollo,
@@ -105,6 +106,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     Lion,
     AliG,
     SM3,
+    AdaNorm,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

{pytorch_optimizer-2.6.0 → pytorch_optimizer-2.7.0}/pytorch_optimizer/optimizer/adabelief.py

@@ -21,7 +21,9 @@ class AdaBelief(Optimizer, BaseOptimizer):
     :param rectify: bool. perform the rectified update similar to RAdam.
     :param degenerated_to_sgd: bool. perform SGD update when variance of gradient is high.
     :param amsgrad: bool. whether to use the AMSBound variant.
-    :param adamd_debias_term: bool. Only correct the denominator to avoid inflating step sizes early in training.
+    :param r: float. EMA factor. between 0.9 ~ 0.99 is preferred.
+    :param adanorm: bool. whether to use the AdaNorm variant.
+    :param adam_debias: bool. Only correct the denominator to avoid inflating step sizes early in training.
     :param eps: float. term added to the denominator to improve numerical stability.
     """
 
@@ -37,18 +39,16 @@ class AdaBelief(Optimizer, BaseOptimizer):
         rectify: bool = True,
         degenerated_to_sgd: bool = True,
         amsgrad: bool = False,
-        adamd_debias_term: bool = False,
+        r: float = 0.95,
+        adanorm: bool = False,
+        adam_debias: bool = False,
         eps: float = 1e-16,
     ):
         self.lr = lr
         self.betas = betas
         self.weight_decay = weight_decay
         self.n_sma_threshold = n_sma_threshold
-        self.weight_decouple = weight_decouple
-        self.fixed_decay = fixed_decay
-        self.rectify = rectify
         self.degenerated_to_sgd = degenerated_to_sgd
-        self.adamd_debias_term = adamd_debias_term
         self.eps = eps
 
         self.validate_parameters()
@@ -56,12 +56,18 @@ class AdaBelief(Optimizer, BaseOptimizer):
         defaults: DEFAULTS = {
             'lr': lr,
             'betas': betas,
-            'eps': eps,
             'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'rectify': rectify,
             'amsgrad': amsgrad,
-            'adamd_debias_term': adamd_debias_term,
-            'buffer': [[None, None, None] for _ in range(10)],
+            'adanorm': adanorm,
+            'adam_debias': adam_debias,
+            'eps': eps,
         }
+        if adanorm:
+            defaults.update({'r': r})
+
         super().__init__(params, defaults)
 
     def validate_parameters(self):
@@ -76,12 +82,14 @@ class AdaBelief(Optimizer, BaseOptimizer):
     @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['step'] = 0
                 state['exp_avg'] = torch.zeros_like(p)
                 state['exp_avg_var'] = torch.zeros_like(p)
+                if group['adanorm']:
+                    state['exp_grad_norm'] = torch.zeros((1,), dtype=p.dtype, device=p.device)
                 if group['amsgrad']:
                     state['max_exp_avg_var'] = torch.zeros_like(p)
 
@@ -93,11 +101,21 @@ class AdaBelief(Optimizer, BaseOptimizer):
                 loss = closure()
 
         for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
             beta1, beta2 = group['betas']
-            weight_decay: float = group['weight_decay']
+            weight_decay = group['weight_decay']
+
+            bias_correction1 = 1.0 - beta1 ** group['step']
+            bias_correction2_sq = math.sqrt(1.0 - beta2 ** group['step'])
 
-            if self.rectify:
+            if group['rectify']:
                 n_sma_max: float = 2.0 / (1.0 - beta2) - 1.0
+                beta2_t: float = beta2 ** group['step']
+                n_sma: float = n_sma_max - 2 * group['step'] * beta2_t / (1.0 - beta2_t)
 
             for p in group['params']:
                 if p.grad is None:
@@ -109,70 +127,68 @@ class AdaBelief(Optimizer, BaseOptimizer):
 
                 state = self.state[p]
                 if len(state) == 0:
-                    state['step'] = 0
                     state['exp_avg'] = torch.zeros_like(p)
                     state['exp_avg_var'] = torch.zeros_like(p)
+                    if group['adanorm']:
+                        state['exp_grad_norm'] = torch.zeros((1,), dtype=grad.dtype, device=grad.device)
                     if group['amsgrad']:
                         state['max_exp_avg_var'] = torch.zeros_like(p)
 
-                if self.weight_decouple:
-                    p.mul_(1.0 - (group['lr'] * weight_decay if not self.fixed_decay else weight_decay))
+                if group['weight_decouple']:
+                    p.mul_(1.0 - group['weight_decay'] * (1.0 if group['fixed_decay'] else group['lr']))
                 elif weight_decay > 0.0:
                     grad.add_(p, alpha=weight_decay)
 
-                state['step'] += 1
                 exp_avg, exp_avg_var = state['exp_avg'], state['exp_avg_var']
 
-                bias_correction1 = 1.0 - beta1 ** state['step']
-                bias_correction2_sq = math.sqrt(1.0 - beta2 ** state['step'])
+                s_grad = grad
+                if group['adanorm']:
+                    grad_norm = torch.linalg.norm(grad)
 
-                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
-                grad_residual = grad - exp_avg
-                exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1.0 - beta2).add_(group['eps'])
+                    exp_grad_norm = state['exp_grad_norm']
+                    exp_grad_norm.mul_(group['r']).add_(grad_norm, alpha=1.0 - group['r'])
+
+                    if exp_grad_norm > grad_norm:
+                        s_grad *= exp_grad_norm / grad_norm
+
+                exp_avg.mul_(beta1).add_(s_grad, alpha=1.0 - beta1)
+                grad_residual = s_grad - exp_avg
+                exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1.0 - beta2).add_(self.eps)
 
                 if group['amsgrad']:
                     max_exp_avg_var = state['max_exp_avg_var']
                     torch.max(max_exp_avg_var, exp_avg_var, out=max_exp_avg_var)
-                    de_nom = max_exp_avg_var.sqrt()
+                    de_nom = max_exp_avg_var.add(self.eps).sqrt()
                 else:
-                    de_nom = exp_avg_var.sqrt()
-                de_nom.div_(bias_correction2_sq).add_(group['eps'])
+                    de_nom = exp_avg_var.add(self.eps).sqrt()
+
+                de_nom.div_(bias_correction2_sq).add_(self.eps)
 
-                if not self.rectify:
-                    step_size: float = group['lr'] if group['adamd_debias_term'] else group['lr'] / bias_correction1
+                if not group['rectify']:
+                    step_size: float = group['lr'] if group['adam_debias'] else group['lr'] / bias_correction1
                     p.addcdiv_(exp_avg, de_nom, value=-step_size)
                     continue
 
-                buffered = group['buffer'][state['step'] % 10]
-                if state['step'] == buffered[0]:
-                    n_sma, step_size = buffered[1], buffered[2]
+                if n_sma >= self.n_sma_threshold:
+                    step_size = math.sqrt(
+                        (1 - beta2_t)
+                        * (n_sma - 4)
+                        / (n_sma_max - 4)
+                        * (n_sma - 2)
+                        / n_sma
+                        * n_sma_max
+                        / (n_sma_max - 2)
+                    )
+                elif self.degenerated_to_sgd:
+                    step_size = 1.0
                 else:
-                    buffered[0] = state['step']
-                    beta2_t = beta2 ** state['step']
-                    n_sma = n_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
-                    buffered[1] = n_sma
-
-                    if n_sma >= self.n_sma_threshold:
-                        step_size = math.sqrt(
-                            (1 - beta2_t)
-                            * (n_sma - 4)
-                            / (n_sma_max - 4)
-                            * (n_sma - 2)
-                            / n_sma
-                            * n_sma_max
-                            / (n_sma_max - 2)
-                        )
-                        if not group['adamd_debias_term']:
-                            step_size /= bias_correction1
-                    elif self.degenerated_to_sgd:
-                        step_size = 1.0 / bias_correction1
-                    else:
-                        step_size = -1
-
-                    buffered[2] = step_size
+                    step_size = -1
+
+                if not group['adam_debias']:
+                    step_size /= bias_correction1
 
                 if n_sma >= self.n_sma_threshold:
-                    de_nom = exp_avg_var.sqrt().add_(group['eps'])
+                    de_nom = exp_avg_var.sqrt().add_(self.eps)
                     p.addcdiv_(exp_avg, de_nom, value=-step_size * group['lr'])
                 elif step_size > 0:
                     p.add_(exp_avg, alpha=-step_size * group['lr'])

{pytorch_optimizer-2.6.0 → pytorch_optimizer-2.7.0}/pytorch_optimizer/optimizer/adabound.py

@@ -21,7 +21,7 @@ class AdaBound(Optimizer, BaseOptimizer):
     :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
     :param fixed_decay: bool. fix weight decay.
     :param amsbound: bool. whether to use the AMSBound variant.
-    :param adamd_debias_term: bool. Only correct the denominator to avoid inflating step sizes early in training.
+    :param adam_debias_term: bool. Only correct the denominator to avoid inflating step sizes early in training.
     :param eps: float. term added to the denominator to improve numerical stability.
     """
 
@@ -36,14 +36,12 @@ class AdaBound(Optimizer, BaseOptimizer):
         weight_decouple: bool = True,
         fixed_decay: bool = False,
         amsbound: bool = False,
-        adamd_debias_term: bool = False,
+        adam_debias: bool = False,
         eps: float = 1e-8,
     ):
         self.lr = lr
         self.betas = betas
         self.weight_decay = weight_decay
-        self.weight_decouple = weight_decouple
-        self.fixed_decay = fixed_decay
         self.eps = eps
 
         self.validate_parameters()
@@ -54,8 +52,10 @@ class AdaBound(Optimizer, BaseOptimizer):
             'final_lr': final_lr,
             'gamma': gamma,
             'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
             'amsbound': amsbound,
-            'adamd_debias_term': adamd_debias_term,
+            'adam_debias': adam_debias,
             'eps': eps,
         }
         super().__init__(params, defaults)
@@ -91,7 +91,20 @@ class AdaBound(Optimizer, BaseOptimizer):
                 loss = closure()
 
         for group, base_lr in zip(self.param_groups, self.base_lrs):
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
             beta1, beta2 = group['betas']
+
+            bias_correction1 = 1.0 - beta1 ** group['step']
+            bias_correction2_sq = math.sqrt(1.0 - beta2 ** group['step'])
+
+            final_lr = group['final_lr'] * group['lr'] / base_lr
+            lower_bound = final_lr * (1 - 1 / (group['gamma'] * group['step'] + 1))
+            upper_bound = final_lr * (1 + 1 / (group['gamma'] * group['step']))
+
             for p in group['params']:
                 if p.grad is None:
                     continue
@@ -103,41 +116,32 @@ class AdaBound(Optimizer, BaseOptimizer):
                 state = self.state[p]
 
                 if len(state) == 0:
-                    state['step'] = 0
                     state['exp_avg'] = torch.zeros_like(p)
                     state['exp_avg_sq'] = torch.zeros_like(p)
                     if group['amsbound']:
                         state['max_exp_avg_sq'] = torch.zeros_like(p)
 
-                state['step'] += 1
-                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                if group['weight_decouple']:
+                    p.mul_(1.0 - group['weight_decay'] * (1.0 if group['fixed_decay'] else group['lr']))
+                elif group['weight_decay'] > 0.0:
+                    grad.add_(p, alpha=group['weight_decay'])
 
-                if group['weight_decay'] > 0.0:
-                    if self.weight_decouple:
-                        p.mul_(
-                            1.0 - (group['weight_decay'] if self.fixed_decay else group['lr'] * group['weight_decay'])
-                        )
-                    else:
-                        grad.add_(p, alpha=group['weight_decay'])
+                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)
-                if group['amsbound']:
-                    torch.max(state['max_exp_avg_sq'], exp_avg_sq, out=exp_avg_sq)
-
-                de_nom = exp_avg_sq.sqrt().add_(group['eps'])
 
-                bias_correction1 = 1.0 - beta1 ** state['step']
-                bias_correction2_sq = math.sqrt(1.0 - beta2 ** state['step'])
+                if group['amsbound']:
+                    max_exp_avg_sq = state['max_exp_avg_sq']
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    de_nom = max_exp_avg_sq.add(group['eps']).sqrt()
+                else:
+                    de_nom = exp_avg_sq.add(group['eps']).sqrt()
 
                 step_size = group['lr'] * bias_correction2_sq
-                if not group['adamd_debias_term']:
+                if not group['adam_debias']:
                     step_size /= bias_correction1
 
-                final_lr = group['final_lr'] * group['lr'] / base_lr
-                lower_bound = final_lr * (1 - 1 / (group['gamma'] * state['step'] + 1))
-                upper_bound = final_lr * (1 + 1 / (group['gamma'] * state['step']))
-
                 step_size = torch.full_like(de_nom, step_size)
                 step_size.div_(de_nom).clamp_(lower_bound, upper_bound).mul_(exp_avg)
 

{pytorch_optimizer-2.6.0 → pytorch_optimizer-2.7.0}/pytorch_optimizer/optimizer/adafactor.py

@@ -18,10 +18,11 @@ class AdaFactor(Optimizer, BaseOptimizer):
     :param decay_rate: float. coefficient used to compute running averages of square gradient.
     :param weight_decay: float. weight decay (L2 penalty).
     :param clip_threshold: float. threshold of root-mean-square of final gradient update.
+    :param amsgrad: bool. whether to use the AMSBound variant.
     :param scale_parameter: bool. if true, learning rate is scaled by root-mean-square of parameter.
     :param relative_step: bool. if true, time-dependent learning rate is computed instead of external learning rate.
-    :param warmup_init: bool. time-dependent learning rate computation depends on
-        whether warm-up initialization is being used.
+    :param warmup_init: bool. time-dependent learning rate computation depends on whether warm-up initialization
+        is being used.
     :param eps1: float. term added to the denominator to improve numerical stability.
     :param eps2: float. term added to the denominator to improve numerical stability.
     """
@@ -34,6 +35,7 @@ class AdaFactor(Optimizer, BaseOptimizer):
         decay_rate: float = -0.8,
         weight_decay: float = 0.0,
         clip_threshold: float = 1.0,
+        amsgrad: bool = False,
         scale_parameter: bool = True,
         relative_step: bool = True,
         warmup_init: bool = False,
@@ -45,15 +47,23 @@ class AdaFactor(Optimizer, BaseOptimizer):
         self.decay_rate = decay_rate
         self.weight_decay = weight_decay
         self.clip_threshold = clip_threshold
-        self.scale_parameter = scale_parameter
+        self.amsgrad = amsgrad
         self.relative_step = relative_step
-        self.warmup_init = warmup_init
         self.eps1 = eps1
         self.eps2 = eps2
 
         self.validate_parameters()
 
-        defaults: DEFAULTS = {'lr': lr, 'weight_decay': weight_decay}
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'weight_decay': weight_decay,
+            'amsgrad': amsgrad,
+            'scale_parameter': scale_parameter,
+            'relative_step': relative_step,
+            'warmup_init': warmup_init,
+            'eps1': eps1,
+            'eps2': eps2,
+        }
         super().__init__(params, defaults)
 
     def validate_parameters(self):
@@ -69,6 +79,7 @@ class AdaFactor(Optimizer, BaseOptimizer):
     @torch.no_grad()
     def reset(self):
         for group in self.param_groups:
+            group['step'] = 0
             for p in group['params']:
                 state = self.state[p]
 
@@ -77,24 +88,31 @@ class AdaFactor(Optimizer, BaseOptimizer):
                 grad_shape: Tuple[int, ...] = grad.shape
                 factored: bool = self.get_options(grad_shape)
 
-                state['step'] = 0
                 state['exp_avg'] = torch.zeros_like(p)
 
                 if factored:
-                    state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1], dtype=grad.dtype)
-                    state['exp_avg_sq_col'] = torch.zeros(grad_shape[:-2] + grad_shape[-1:], dtype=grad.dtype)
+                    state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1], dtype=grad.dtype, device=grad.device)
+                    state['exp_avg_sq_col'] = torch.zeros(
+                        grad_shape[:-2] + grad_shape[-1:], dtype=grad.dtype, device=grad.device
+                    )
                 else:
                     state['exp_avg_sq'] = torch.zeros_like(grad)
 
+                if group['amsgrad']:
+                    state['exp_avg_sq_hat'] = torch.zeros_like(grad)
+
                 state['RMS'] = 0.0
 
-    def get_lr(self, lr: float, step: int, rms: float) -> float:
+    def get_lr(
+        self, lr: float, step: int, rms: float, relative_step: bool, warmup_init: bool, scale_parameter: bool
+    ) -> float:
+        r"""Get AdaFactor learning rate."""
         relative_step_size: float = lr
-        if self.relative_step:
-            min_step: float = 1e-6 * step if self.warmup_init else 1e-2
+        if relative_step:
+            min_step: float = 1e-6 * step if warmup_init else 1e-2
             relative_step_size = min(min_step, 1.0 / math.sqrt(step))
 
-        param_scale: float = 1.0 if self.scale_parameter else max(self.eps2, rms)
+        param_scale: float = 1.0 if scale_parameter else max(self.eps2, rms)
 
         return param_scale * relative_step_size
 
@@ -114,8 +132,8 @@ class AdaFactor(Optimizer, BaseOptimizer):
         exp_avg_sq_col: torch.Tensor,
         output: torch.Tensor,
     ):
-        r"""Get approximate squared gradient."""
-        r_factor: torch.Tensor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1)).rsqrt_().unsqueeze(-1)
+        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)
 
@@ -127,6 +145,13 @@ class AdaFactor(Optimizer, BaseOptimizer):
                 loss = closure()
 
         for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            beta2_t: float = 1.0 - math.pow(group['step'], self.decay_rate)
+
             for p in group['params']:
                 if p.grad is None:
                     continue
@@ -141,23 +166,32 @@ class AdaFactor(Optimizer, BaseOptimizer):
                 factored: bool = self.get_options(grad_shape)
 
                 if len(state) == 0:
-                    state['step'] = 0
                     state['exp_avg'] = torch.zeros_like(p)
 
                     if factored:
-                        state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1], dtype=grad.dtype)
-                        state['exp_avg_sq_col'] = torch.zeros(grad_shape[:-2] + grad_shape[-1:], dtype=grad.dtype)
+                        state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1], dtype=grad.dtype, device=grad.device)
+                        state['exp_avg_sq_col'] = torch.zeros(
+                            grad_shape[:-2] + grad_shape[-1:], dtype=grad.dtype, device=grad.device
+                        )
                     else:
                         state['exp_avg_sq'] = torch.zeros_like(grad)
 
-                    state['RMS'] = 0.0
+                    if group['amsgrad']:
+                        state['exp_avg_sq_hat'] = torch.zeros_like(grad)
 
-                state['step'] += 1
+                    state['RMS'] = 0.0
 
                 state['RMS'] = self.get_rms(p)
-                lr = self.get_lr(group['lr'], state['step'], state['RMS'])
 
-                beta2_t: float = 1.0 - math.pow(state['step'], self.decay_rate)
+                lr = self.get_lr(
+                    group['lr'],
+                    group['step'],
+                    state['RMS'],
+                    relative_step=group['relative_step'],
+                    warmup_init=group['warmup_init'],
+                    scale_parameter=group['scale_parameter'],
+                )
+
                 update = torch.mul(grad, grad).add_(self.eps1)
 
                 if factored:
@@ -166,28 +200,27 @@ class AdaFactor(Optimizer, BaseOptimizer):
                     exp_avg_sq_row.mul_(beta2_t).add_(update.mean(dim=-1), alpha=1.0 - beta2_t)
                     exp_avg_sq_col.mul_(beta2_t).add_(update.mean(dim=-2), alpha=1.0 - beta2_t)
 
-                    # Approximation of exponential moving average of square of gradient
                     self.approximate_sq_grad(exp_avg_sq_row, exp_avg_sq_col, update)
                 else:
                     exp_avg_sq = state['exp_avg_sq']
-
                     exp_avg_sq.mul_(beta2_t).add_(update, alpha=1.0 - beta2_t)
                     torch.rsqrt(exp_avg_sq, out=update)
 
-                update.mul_(grad)
+                if group['amsgrad']:
+                    exp_avg_sq_hat = state['exp_avg_sq_hat']
+                    torch.max(exp_avg_sq_hat, 1 / update, out=exp_avg_sq_hat)
+                    torch.rsqrt(exp_avg_sq_hat / beta2_t, out=update)
 
-                # TODO: implement AMSGrad
+                update.mul_(grad)
 
                 update.div_((self.get_rms(update) / self.clip_threshold).clamp_(min=1.0)).mul_(lr)
 
                 exp_avg = state['exp_avg']
                 exp_avg.mul_(self.betas[0]).add_(update, alpha=1.0 - self.betas[0])
 
-                update = exp_avg
-
                 if group['weight_decay'] > 0.0:
                     p.add_(p, alpha=-lr * group['weight_decay'])
 
-                p.add_(-update)
+                p.add_(-exp_avg)
 
         return loss