Repository: kozistr/pytorch_optimizer

Issue:
Optimizer Wrapper Inconsistency: OrthoGrad, Lookahead, and TRAC – Streamlining Optimizer Passing
There’s an inconsistency in how optimizer wrappers like OrthoGrad, Lookahead, and TRAC handle the base optimizer:

- OrthoGrad lets you pass the optimizer class directly or create it inside:
```python
optimizer = OrthoGrad(params, optimizer=torch.optim.AdamW, lr=1e-3)  
```
- Lookahead and TRAC require a fully instantiated optimizer object:
```python
base_optimizer = torch.optim.AdamW(params, lr=1e-3)  
optimizer = Lookahead(optimizer=base_optimizer, k=5, alpha=0.5)  
```

This means you can’t seamlessly chain optimizers like:

```python
optimizer = OrthoGrad(params, optimizer=Lookahead(torch.optim.AdamW, k=5), lr=1e-3) 
```

If you want to compose them, you’re stuck figuring out what each one needs, and that makes the API clunky and less user-friendly.

Proposed Fix:

Allow passing either an optimizer class (like torch.optim.AdamW) with its params or an already-created optimizer instance.
Use a consistent interface that takes care of the internal setup automatically.

```python

class Lookahead(BaseOptimizer):  
    def __init__(  
        self,  
        params: PARAMETERS,  
        optimizer: Union[OPTIMIZER, Type[OPTIMIZER]] = torch.optim.AdamW,  
        **kwargs  
    ):  
        # If optimizer is a class, instantiate it  
        if isinstance(optimizer, type):  
            self.optimizer = optimizer(params, **kwargs)  
        else:  
            self.optimizer = optimizer  

        # Rest of the initialization  
        self.alpha = kwargs.pop('alpha', 0.5)  
        self.k = kwargs.pop('k', 5)  
        # ...  
```

This keeps backward compatibility while making the API consistent and composable.

Things like this could be possible:

```python
optimizer = OrthoGrad(  
    params,  
    optimizer=Lookahead(  
        optimizer=torch.optim.AdamW,  
        k=5,  
        alpha=0.5  
    ),  
    lr=1e-3  
)  
```

Hints from issue comments:
@kozistr: Thanks for your suggestion! Supporting either an optimizer class with its params or an already-created optimizer instance also seems great. I'll work on it :)

thanks!

Solution patch:
diff --git a/README.md b/README.md
index 687d919a..d252b0f4 100644
--- a/README.md
+++ b/README.md
@@ -10,8 +10,8 @@
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **89 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
-* Including many variants such as `Cautious`, `AdamD`, `Gradient Centrailiaztion`
+* Wide range of supported optimizers. Currently, **90 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Including many variants such as `ADOPT`, `Cautious`, `AdamD`, `StableAdamW`, and `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
 * Somewhat a bit more optimized compared to the original implementation
@@ -198,6 +198,7 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | Grams         | *Gradient Descent with Adaptive Momentum Scaling*                                                 |                                                                                                                | <https://arxiv.org/abs/2412.17107>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv241217107C/exportcitation)                                                        |
 | OrthoGrad     | *Grokking at the Edge of Numerical Stability*                                                     | [github](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)                         | <https://arxiv.org/abs/2501.04697>                                                          | [cite](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability?tab=readme-ov-file#citation)                    |
 | Adam-ATAN2    | *Scaling Exponents Across Parameterizations and Optimizers*                                       |                                                                                                                | <https://arxiv.org/abs/2407.05872>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240705872E/exportcitation)                                                        |
+| SPAM          | *Spike-Aware Adam with Momentum Reset for Stable LLM Training*                                    | [github](https://github.com/TianjinYellow/SPAM-Optimizer)                                                      | <https://arxiv.org/abs/2501.06842>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250106842H/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 
diff --git a/docs/changelogs/v3.3.4.md b/docs/changelogs/v3.3.4.md
new file mode 100644
index 00000000..7e6f27be
--- /dev/null
+++ b/docs/changelogs/v3.3.4.md
@@ -0,0 +1,12 @@
+### Change Log
+
+### Feature
+
+* Support `OrthoGrad` feature for `create_optimizer()`. (#324)
+* Enhanced flexibility for the `optimizer` parameter in `Lookahead`, `TRAC`, and `OrthoGrad` optimizers. (#324)
+    * Now supports both torch.optim.Optimizer instances and classes
+    * You can now use `Lookahead` optimizer in two ways.
+        * `Lookahead(AdamW(model.parameters(), lr=1e-3), k=5, alpha=0.5)`
+        * `Lookahead(AdamW, k=5, alpha=0.5, params=model.parameters())`
+* Implement `SPAM` optimizer. (#324)
+    * [Spike-Aware Adam with Momentum Reset for Stable LLM Training](https://arxiv.org/abs/2501.06842)
diff --git a/docs/index.md b/docs/index.md
index 687d919a..d252b0f4 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -10,8 +10,8 @@
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **89 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
-* Including many variants such as `Cautious`, `AdamD`, `Gradient Centrailiaztion`
+* Wide range of supported optimizers. Currently, **90 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Including many variants such as `ADOPT`, `Cautious`, `AdamD`, `StableAdamW`, and `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
 * Somewhat a bit more optimized compared to the original implementation
@@ -198,6 +198,7 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | Grams         | *Gradient Descent with Adaptive Momentum Scaling*                                                 |                                                                                                                | <https://arxiv.org/abs/2412.17107>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv241217107C/exportcitation)                                                        |
 | OrthoGrad     | *Grokking at the Edge of Numerical Stability*                                                     | [github](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)                         | <https://arxiv.org/abs/2501.04697>                                                          | [cite](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability?tab=readme-ov-file#citation)                    |
 | Adam-ATAN2    | *Scaling Exponents Across Parameterizations and Optimizers*                                       |                                                                                                                | <https://arxiv.org/abs/2407.05872>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240705872E/exportcitation)                                                        |
+| SPAM          | *Spike-Aware Adam with Momentum Reset for Stable LLM Training*                                    | [github](https://github.com/TianjinYellow/SPAM-Optimizer)                                                      | <https://arxiv.org/abs/2501.06842>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250106842H/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 
diff --git a/docs/optimizer.md b/docs/optimizer.md
index 7f868612..a5a56088 100644
--- a/docs/optimizer.md
+++ b/docs/optimizer.md
@@ -368,6 +368,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.SPAM
+    :docstring:
+    :members:
+
 ::: pytorch_optimizer.SRMM
     :docstring:
     :members:
diff --git a/pyproject.toml b/pyproject.toml
index 6d602015..3d83a284 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,9 +18,9 @@ keywords = [
     "Kate", "Lamb", "LaProp", "LARS", "Lion", "LOMO", "Lookahead", "MADGRAD", "MARS", "MSVAG", "Muno", "Nero",
     "NovoGrad", "OrthoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger",
     "Ranger21", "RotoGrad", "SAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SGDP", "Shampoo",
-    "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH", "SRMM", "StableAdamW", "SWATS", "Tiger", "TRAC",
-    "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered",
-    "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
+    "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH", "SPAM", "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",
diff --git a/pytorch_optimizer/__init__.py b/pytorch_optimizer/__init__.py
index b5fb37b9..66be5d6c 100644
--- a/pytorch_optimizer/__init__.py
+++ b/pytorch_optimizer/__init__.py
@@ -61,6 +61,7 @@
     SGDW,
     SM3,
     SOAP,
+    SPAM,
     SRMM,
     SWATS,
     TRAC,
diff --git a/pytorch_optimizer/base/scheduler.py b/pytorch_optimizer/base/scheduler.py
index 4cfd5ed3..11c78f72 100644
--- a/pytorch_optimizer/base/scheduler.py
+++ b/pytorch_optimizer/base/scheduler.py
@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
 from typing import List
 
+from torch.optim import Optimizer
+
 from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError
-from pytorch_optimizer.base.types import OPTIMIZER
 
 
 class BaseLinearWarmupScheduler(ABC):
@@ -10,7 +11,7 @@ class BaseLinearWarmupScheduler(ABC):
 
         The LR Scheduler class based on this class has linear warmup strategy.
 
-    :param optimizer: Optimizer. OPTIMIZER. It will set learning rate to all trainable parameters in optimizer.
+    :param optimizer: Optimizer. It will set learning rate to all trainable parameters in optimizer.
     :param t_max: int. total steps to train.
     :param max_lr: float. maximum lr.
     :param min_lr: float. minimum lr.
@@ -20,7 +21,7 @@ class BaseLinearWarmupScheduler(ABC):
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: Optimizer,
         t_max: int,
         max_lr: float,
         min_lr: float = 0.0,
diff --git a/pytorch_optimizer/base/types.py b/pytorch_optimizer/base/types.py
index ff6aaaf9..786a9c08 100644
--- a/pytorch_optimizer/base/types.py
+++ b/pytorch_optimizer/base/types.py
@@ -11,6 +11,7 @@
 PARAMETERS = Optional[Union[Iterable[Dict], Iterable[torch.Tensor]]]
 STATE = Dict
 OPTIMIZER = Type[Optimizer]
+OPTIMIZER_INSTANCE_OR_CLASS = Union[OPTIMIZER, Optimizer]
 SCHEDULER = Type[LRScheduler]
 
 HUTCHINSON_G = Literal['gaussian', 'rademacher']
diff --git a/pytorch_optimizer/lr_scheduler/cosine_anealing.py b/pytorch_optimizer/lr_scheduler/cosine_anealing.py
index 3d2102af..42e282e9 100644
--- a/pytorch_optimizer/lr_scheduler/cosine_anealing.py
+++ b/pytorch_optimizer/lr_scheduler/cosine_anealing.py
@@ -1,10 +1,9 @@
 import math
 from typing import List, Optional
 
+from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
-from pytorch_optimizer.base.types import OPTIMIZER
-
 
 class CosineAnnealingWarmupRestarts(LRScheduler):
     r"""CosineAnnealingWarmupRestarts.
@@ -21,7 +20,7 @@ class CosineAnnealingWarmupRestarts(LRScheduler):
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: Optimizer,
         first_cycle_steps: int,
         cycle_mult: float = 1.0,
         max_lr: float = 1e-4,
@@ -53,7 +52,7 @@ def __init__(
 
         self.init_lr()
 
-    def init_lr(self):
+    def init_lr(self) -> None:
         self.base_lrs = []
         for param_group in self.optimizer.param_groups:
             param_group['lr'] = self.min_lr
diff --git a/pytorch_optimizer/lr_scheduler/rex.py b/pytorch_optimizer/lr_scheduler/rex.py
index 5d7ad3d3..9dda5ac2 100644
--- a/pytorch_optimizer/lr_scheduler/rex.py
+++ b/pytorch_optimizer/lr_scheduler/rex.py
@@ -1,9 +1,8 @@
 from typing import List, Optional
 
+from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
-from pytorch_optimizer.base.types import OPTIMIZER
-
 
 class REXScheduler(LRScheduler):
     r"""Revisiting Budgeted Training with an Improved Schedule.
@@ -16,7 +15,7 @@ class REXScheduler(LRScheduler):
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: Optimizer,
         total_steps: int,
         max_lr: float = 1.0,
         min_lr: float = 0.0,
@@ -35,7 +34,7 @@ def __init__(
 
         self.init_lr()
 
-    def init_lr(self):
+    def init_lr(self) -> None:
         self.base_lrs = []
         for param_group in self.optimizer.param_groups:
             param_group['lr'] = self.min_lr
diff --git a/pytorch_optimizer/optimizer/__init__.py b/pytorch_optimizer/optimizer/__init__.py
index bde4ec5a..c80c459e 100644
--- a/pytorch_optimizer/optimizer/__init__.py
+++ b/pytorch_optimizer/optimizer/__init__.py
@@ -1,10 +1,10 @@
 import fnmatch
 from importlib.util import find_spec
-from typing import Dict, List, Optional, Sequence, Set, Union
+from typing import Dict, List, Optional, Sequence, Set, Type, Union
 
 import torch
 from torch import nn
-from torch.optim import AdamW
+from torch.optim import AdamW, Optimizer
 
 from pytorch_optimizer.base.types import OPTIMIZER, PARAMETERS
 from pytorch_optimizer.optimizer.a2grad import A2Grad
@@ -83,6 +83,7 @@
 from pytorch_optimizer.optimizer.sm3 import SM3
 from pytorch_optimizer.optimizer.soap import SOAP
 from pytorch_optimizer.optimizer.sophia import SophiaH
+from pytorch_optimizer.optimizer.spam import SPAM
 from pytorch_optimizer.optimizer.srmm import SRMM
 from pytorch_optimizer.optimizer.swats import SWATS
 from pytorch_optimizer.optimizer.tiger import Tiger
@@ -286,6 +287,7 @@ def load_optimizer(optimizer: str) -> OPTIMIZER:
     MARS,
     SGDSaI,
     Grams,
+    SPAM,
     Ranger25,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
@@ -298,8 +300,9 @@ def create_optimizer(
     weight_decay: float = 0.0,
     wd_ban_list: List[str] = ('bias', 'LayerNorm.bias', 'LayerNorm.weight'),
     use_lookahead: bool = False,
+    use_orthograd: bool = False,
     **kwargs,
-):
+) -> Optimizer:
     r"""Build optimizer.
 
     :param model: nn.Module. model.
@@ -307,7 +310,8 @@ def create_optimizer(
     :param lr: float. learning rate.
     :param weight_decay: float. weight decay.
     :param wd_ban_list: List[str]. weight decay ban list by layer.
-    :param use_lookahead: bool. use lookahead.
+    :param use_lookahead: bool. use Lookahead.
+    :param use_orthograd: bool. use OrthoGrad.
     """
     optimizer_name = optimizer_name.lower()
 
@@ -315,14 +319,17 @@ def create_optimizer(
         get_optimizer_parameters(model, weight_decay, wd_ban_list) if weight_decay > 0.0 else model.parameters()
     )
 
-    optimizer = load_optimizer(optimizer_name)
+    optimizer_class: OPTIMIZER = load_optimizer(optimizer_name)
 
     if optimizer_name == 'alig':
-        optimizer = optimizer(parameters, max_lr=lr, **kwargs)
+        optimizer = optimizer_class(parameters, max_lr=lr, **kwargs)
     elif optimizer_name in {'lomo', 'adalomo', 'adammini'}:
-        optimizer = optimizer(model, lr=lr, **kwargs)
+        optimizer = optimizer_class(model, lr=lr, **kwargs)
     else:
-        optimizer = optimizer(parameters, lr=lr, **kwargs)
+        optimizer = optimizer_class(parameters, lr=lr, **kwargs)
+
+    if use_orthograd:
+        optimizer = OrthoGrad(optimizer, **kwargs)
 
     if use_lookahead:
         optimizer = Lookahead(
diff --git a/pytorch_optimizer/optimizer/experimental/ranger25.py b/pytorch_optimizer/optimizer/experimental/ranger25.py
index 6de13e3e..f30f2340 100644
--- a/pytorch_optimizer/optimizer/experimental/ranger25.py
+++ b/pytorch_optimizer/optimizer/experimental/ranger25.py
@@ -11,6 +11,8 @@
 class Ranger25(BaseOptimizer):
     r"""Mixin' every fancy optimizer hacks.
 
+    ADOPT + AdEMAMix + Cautious + StableAdamW + Adam-Atan2
+
     :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.
@@ -19,10 +21,10 @@ class Ranger25(BaseOptimizer):
     :param fixed_decay: bool. fix weight decay.
     :param alpha: float. usually between 4 and 10 would work well.
     :param t_alpha_beta3: Optional[float]. total number of iterations is preferred when needed.
-    :param n_sma_threshold: number of SMA threshold (recommended is 5).
     :param cautious: bool. whether to use the Cautious variant.
     :param stable_adamw: bool. whether to use stable AdamW variant.
-    :param eps: float. term added to the denominator to improve numerical stability.
+    :param eps: Optional[float]. term added to the denominator to improve numerical stability. when eps is None and
+        stable_adamw is False, adam-atan2 feature will be used.
     """
 
     def __init__(
@@ -35,10 +37,9 @@ def __init__(
         fixed_decay: bool = False,
         alpha: float = 5.0,
         t_alpha_beta3: Optional[float] = None,
-        n_sma_threshold: int = 5,
         cautious: bool = True,
         stable_adamw: bool = True,
-        eps: float = 1e-8,
+        eps: Optional[float] = 1e-8,
         **kwargs,
     ):
         self.validate_learning_rate(lr)
@@ -48,9 +49,8 @@ def __init__(
         self.validate_non_negative(weight_decay, 'weight_decay')
         self.validate_non_negative(eps, 'eps')
 
-        self.n_sma_threshold = n_sma_threshold
         self.cautious = cautious
-        self.stable_adamw = stable_adamw
+        self.stable_adamw: bool = stable_adamw if isinstance(eps, float) else False
 
         defaults: DEFAULTS = {
             'lr': lr,
@@ -60,7 +60,7 @@ def __init__(
             'fixed_decay': fixed_decay,
             'alpha': alpha,
             't_alpha_beta3': t_alpha_beta3,
-            'eps': eps,
+            'eps': eps if (eps is not None) or (eps is None and not stable_adamw) else 1e-8,
         }
 
         super().__init__(params, defaults)
@@ -147,16 +147,14 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 exp_avg, exp_avg_sq, exp_avg_slow = state['exp_avg'], state['exp_avg_sq'], state['exp_avg_slow']
 
-                de_nom = exp_avg_sq.sqrt().clamp_(min=group['eps'])
-
-                normed_grad = grad.div(de_nom).clamp_(-clip, clip)
+                normed_grad = grad.div(
+                    exp_avg_sq.sqrt().clamp_(min=group['eps'] if group['eps'] is not None else 1e-8)
+                ).clamp_(-clip, clip)
 
                 exp_avg.mul_(beta1).add_(normed_grad, alpha=1.0 - beta1)
                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                 exp_avg_slow.mul_(beta3_t).add_(normed_grad, alpha=1.0 - beta3_t)
 
-                de_nom = exp_avg_sq.sqrt().div_(bias_correction2_sq).add_(group['eps'])
-
                 update = exp_avg.clone()
                 if self.cautious:
                     self.apply_cautious(update, grad)
@@ -164,21 +162,14 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 if self.stable_adamw:
                     step_size /= self.get_stable_adamw_rms(grad, exp_avg_sq)
 
-                step_size, n_sma = self.get_rectify_step_size(
-                    is_rectify=True,
-                    step=group['step'],
-                    lr=step_size,
-                    beta2=beta2,
-                    n_sma_threshold=self.n_sma_threshold,
-                    degenerated_to_sgd=False,
-                )
+                update.add_(exp_avg_slow, alpha=alpha_t)
+
+                de_nom = exp_avg_sq.sqrt().div_(bias_correction2_sq)
 
-                update.add_(exp_avg_slow, alpha=alpha_t).div_(de_nom)
+                if group['eps'] is not None:
+                    p.addcdiv_(update, de_nom.add_(group['eps']), value=-step_size)
+                    continue
 
-                if n_sma >= self.n_sma_threshold:
-                    de_nom = exp_avg_sq.sqrt().add_(group['eps'])
-                    p.addcdiv_(update, de_nom, value=-step_size)
-                else:
-                    p.add_(update, alpha=-step_size)
+                p.add_(update.atan2_(de_nom), alpha=-step_size)
 
         return loss
diff --git a/pytorch_optimizer/optimizer/fp16.py b/pytorch_optimizer/optimizer/fp16.py
index fc7c72b1..f40d15f9 100644
--- a/pytorch_optimizer/optimizer/fp16.py
+++ b/pytorch_optimizer/optimizer/fp16.py
@@ -4,7 +4,7 @@
 from torch import nn
 from torch.optim import Optimizer
 
-from pytorch_optimizer.base.types import CLOSURE, OPTIMIZER, PARAMETERS
+from pytorch_optimizer.base.types import CLOSURE, PARAMETERS
 from pytorch_optimizer.optimizer.utils import clip_grad_norm, has_overflow
 
 
@@ -100,7 +100,7 @@ class SafeFP16Optimizer(Optimizer):  # pragma: no cover
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: Optimizer,
         aggregate_g_norms: bool = False,
         min_loss_scale: float = 2 ** -5,
     ):  # fmt: skip
@@ -113,10 +113,9 @@ def __init__(
 
         # we want the optimizer to be tracking the fp32 parameters
         if len(optimizer.param_groups) != 1:
-            # future implementers: this should hopefully be a matter of just
-            # iterating through the param groups and keeping track of the pointer
-            # through the fp32_params
-            raise NotImplementedError('[-] Need to implement the parameter group transfer.')
+            # future implementers: this should hopefully be a matter of just iterating through the param groups and
+            # keeping track of the pointer through the fp32_params
+            raise NotImplementedError('Need to implement the parameter group transfer.')
 
         optimizer.param_groups[0]['params'] = self.fp32_params
 
@@ -124,17 +123,16 @@ def __init__(
         self.needs_sync: bool = True
 
     @classmethod
-    def get_parameters(cls, optimizer: OPTIMIZER):
+    def get_parameters(cls, optimizer: Optimizer) -> List:
         params: List = []
-        for pg in optimizer.param_groups:
-            params += list(pg['params'])
+        for group in optimizer.param_groups:
+            params += list(group['params'])
         return params
 
     @classmethod
     def build_fp32_params(
         cls, parameters: PARAMETERS, flatten: bool = True
     ) -> Union[torch.Tensor, List[torch.Tensor]]:
-        # create FP32 copy of parameters and grads
         if flatten:
             total_param_size: int = sum(p.numel() for p in parameters)
             fp32_params = torch.zeros(total_param_size, dtype=torch.float, device=parameters[0].device)
@@ -142,7 +140,7 @@ def build_fp32_params(
             offset: int = 0
             for p in parameters:
                 p_num_el = p.numel()
-                fp32_params[offset : offset + p_num_el].copy_(p.view(-1))
+                fp32_params[offset:offset + p_num_el].copy_(p.view(-1))  # fmt: skip
                 offset += p_num_el
 
             fp32_params = nn.Parameter(fp32_params)
@@ -168,9 +166,9 @@ def state_dict(self) -> Dict:
     def load_state_dict(self, state_dict: Dict):
         r"""Load an optimizer state dict.
 
-            In general, we should prefer the configuration of the existing optimizer instance
-            (e.g., learning rate) over that found in the state_dict. This allows us to
-            resume training from a checkpoint using a new set of optimizer args.
+            In general, we should prefer the configuration of the existing optimizer instance (e.g., learning rate)
+            over that found in the state_dict. This allows us to resume training from a checkpoint using a new set of
+            optimizer args.
 
         :param state_dict: Dict. state_dict.
         """
@@ -196,35 +194,34 @@ def backward(self, loss, update_main_grads: bool = False):
         if update_main_grads:
             self.update_main_grads()
 
-    def sync_fp16_grads_to_fp32(self, multiply_grads: float = 1.0):
+    def sync_fp16_grads_to_fp32(self, multiply_grads: float = 1.0) -> None:
         r"""Sync fp16 to fp32 gradients."""
         if self.needs_sync:
             if self.scaler is not None:
-                # correct for dynamic loss scaler
                 multiply_grads /= self.scaler.loss_scale
 
-            # copy FP16 grads to FP32
-            for p, p32 in zip(self.fp16_params, self.fp32_params):
-                if not p.requires_grad:
+            for p16, p32 in zip(self.fp16_params, self.fp32_params):
+                if not p16.requires_grad:
                     continue
 
-                if p.grad is not None:
-                    p32.grad.copy_(p.grad)
+                if p16.grad is not None:
+                    p32.grad.copy_(p16.grad)
                     p32.grad.mul_(multiply_grads)
                 else:
-                    p32.grad = torch.zeros_like(p, dtype=torch.float)
+                    p32.grad = torch.zeros_like(p16, dtype=torch.float)
 
             self.needs_sync = False
 
-    def multiply_grads(self, c: float):
+    def multiply_grads(self, c: float) -> None:
         r"""Multiply grads by a constant c."""
         if self.needs_sync:
             self.sync_fp16_grads_to_fp32(c)
-        else:
-            for p32 in self.fp32_params:
-                p32.grad.mul_(c)
+            return
 
-    def update_main_grads(self):
+        for p32 in self.fp32_params:
+            p32.grad.mul_(c)
+
+    def update_main_grads(self) -> None:
         self.sync_fp16_grads_to_fp32()
 
     def clip_main_grads(self, max_norm: float):
@@ -249,8 +246,7 @@ def clip_main_grads(self, max_norm: float):
 
                     raise FloatingPointError(
                         f'Minimum loss scale reached ({self.min_loss_scale}). Your loss is probably exploding. '
-                        'Try lowering the learning rate, using gradient clipping or '
-                        'increasing the batch size.\n'
+                        'Try lowering the learning rate, using gradient clipping or increasing the batch size.\n'
                         f'Overflow: setting loss scale to {self.scaler.loss_scale}'
                     )
 
@@ -261,16 +257,15 @@ def step(self, closure: CLOSURE = None):
         self.sync_fp16_grads_to_fp32()
         self.optimizer.step(closure)
 
-        # copy FP32 params back into FP16 model
-        for p, p32 in zip(self.fp16_params, self.fp32_params):
-            if not p.requires_grad:
+        for p16, p32 in zip(self.fp16_params, self.fp32_params):
+            if not p16.requires_grad:
                 continue
-            p.data.copy_(p32)
+            p16.data.copy_(p32)
 
-    def zero_grad(self):
+    def zero_grad(self) -> None:
         r"""Clear the gradients of all optimized parameters."""
-        for p in self.fp16_params:
-            p.grad = None
+        for p16 in self.fp16_params:
+            p16.grad = None
         for p32 in self.fp32_params:
             p32.grad.zero_()
         self.needs_sync = False
diff --git a/pytorch_optimizer/optimizer/lookahead.py b/pytorch_optimizer/optimizer/lookahead.py
index 84e84520..50c764e6 100644
--- a/pytorch_optimizer/optimizer/lookahead.py
+++ b/pytorch_optimizer/optimizer/lookahead.py
@@ -2,15 +2,16 @@
 from typing import Callable, Dict
 
 import torch
+from torch.optim import Optimizer
 
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, OPTIMIZER, STATE
+from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, OPTIMIZER_INSTANCE_OR_CLASS, STATE
 
 
 class Lookahead(BaseOptimizer):
     r"""k steps forward, 1 step back.
 
-    :param optimizer: OPTIMIZER. base optimizer.
+    :param optimizer: OPTIMIZER_INSTANCE_OR_CLASS. base optimizer.
     :param k: int. number of lookahead steps.
     :param alpha: float. linear interpolation factor.
     :param pullback_momentum: str. change to inner optimizer momentum on interpolation update.
@@ -18,7 +19,7 @@ class Lookahead(BaseOptimizer):
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: OPTIMIZER_INSTANCE_OR_CLASS,
         k: int = 5,
         alpha: float = 0.5,
         pullback_momentum: str = 'none',
@@ -28,10 +29,17 @@ def __init__(
         self.validate_range(alpha, 'alpha', 0.0, 1.0)
         self.validate_options(pullback_momentum, 'pullback_momentum', ['none', 'reset', 'pullback'])
 
+        if isinstance(optimizer, Optimizer):
+            self.optimizer = optimizer
+        elif 'params' in kwargs:
+            params = kwargs.pop('params')
+            self.optimizer = optimizer(params, **kwargs)
+        else:
+            raise ValueError('Need to pass `params` when you pass the torch.optim.Optimizer instance.')
+
         self._optimizer_step_pre_hooks: Dict[int, Callable] = {}
         self._optimizer_step_post_hooks: Dict[int, Callable] = {}
 
-        self.optimizer = optimizer
         self.alpha = alpha
         self.k = k
         self.pullback_momentum = pullback_momentum
@@ -53,7 +61,7 @@ def __init__(
             'lookahead_alpha': alpha,
             'lookahead_k': k,
             'lookahead_pullback_momentum': pullback_momentum,
-            **optimizer.defaults,
+            **self.optimizer.defaults,
         }
 
     @property
diff --git a/pytorch_optimizer/optimizer/nero.py b/pytorch_optimizer/optimizer/nero.py
index 43622788..48e480a0 100644
--- a/pytorch_optimizer/optimizer/nero.py
+++ b/pytorch_optimizer/optimizer/nero.py
@@ -63,7 +63,7 @@ def reset(self):
             for p in group['params']:
                 if group['constraints'] and p.dim() > 1:
                     p.sub_(neuron_mean(p))
-                    p.div_(neuron_norm(p) + group['eps'])
+                    p.div_(neuron_norm(p).add_(group['eps']))
 
                 state = self.state[p]
 
@@ -94,7 +94,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 if len(state) == 0:
                     if group['constraints'] and p.dim() > 1:
                         p.sub_(neuron_mean(p))
-                        p.div_(neuron_norm(p) + group['eps'])
+                        p.div_(neuron_norm(p).add_(group['eps']))
 
                     state['step'] = 0
                     state['exp_avg_sq'] = torch.zeros_like(neuron_norm(p))
@@ -114,10 +114,10 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 grad_normed = grad / ((exp_avg_sq / bias_correction).sqrt_().add_(group['eps']))
                 torch.nan_to_num(grad_normed, nan=0.0, out=grad_normed)
 
-                p.sub_(grad_normed, alpha=group['lr'] * state['scale'])
+                p.add_(grad_normed, alpha=-group['lr'] * state['scale'])
 
                 if group['constraints'] and p.dim() > 1:
                     p.sub_(neuron_mean(p))
-                    p.div_(neuron_norm(p) + group['eps'])
+                    p.div_(neuron_norm(p).add_(group['eps']))
 
         return loss
diff --git a/pytorch_optimizer/optimizer/orthograd.py b/pytorch_optimizer/optimizer/orthograd.py
index e7880b49..620bdbd6 100644
--- a/pytorch_optimizer/optimizer/orthograd.py
+++ b/pytorch_optimizer/optimizer/orthograd.py
@@ -1,7 +1,10 @@
+from typing import Callable, Dict
+
 import torch
+from torch.optim import Optimizer
 
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import CLOSURE, LOSS, OPTIMIZER, PARAMETERS
+from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, OPTIMIZER_INSTANCE_OR_CLASS
 
 
 class OrthoGrad(BaseOptimizer):
@@ -9,19 +12,34 @@ class OrthoGrad(BaseOptimizer):
 
     A wrapper optimizer that projects gradients to be orthogonal to the current parameters before performing an update.
 
-    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
-    :param optimizer: OPTIMIZER. base optimizer.
+    :param optimizer: OPTIMIZER_INSTANCE_OR_CLASS. base optimizer.
     """
 
-    def __init__(self, params: PARAMETERS, optimizer: OPTIMIZER = torch.optim.AdamW, **kwargs):
+    def __init__(self, optimizer: OPTIMIZER_INSTANCE_OR_CLASS, **kwargs) -> None:
+        self._optimizer_step_pre_hooks: Dict[int, Callable] = {}
+        self._optimizer_step_post_hooks: Dict[int, Callable] = {}
         self.eps: float = 1e-30
 
-        super().__init__(params, {})
-        self.base_optimizer = optimizer(self.param_groups, **kwargs)
+        if isinstance(optimizer, Optimizer):
+            self.optimizer = optimizer
+        elif 'params' in kwargs:
+            params = kwargs.pop('params')
+            self.optimizer = optimizer(params, **kwargs)
+        else:
+            raise ValueError('Need to pass `params` when you pass the torch.optim.Optimizer instance.')
+
+        self.defaults: DEFAULTS = self.optimizer.defaults
 
     def __str__(self) -> str:
         return 'OrthoGrad'
 
+    @property
+    def param_groups(self):
+        return self.optimizer.param_groups
+
+    def __getstate__(self):
+        return {'optimizer': self.optimizer}
+
     @torch.no_grad()
     def reset(self):
         pass
@@ -45,4 +63,4 @@ def orthogonalize_gradients(self, params) -> None:
     def step(self, closure: CLOSURE = None) -> LOSS:
         for group in self.param_groups:
             self.orthogonalize_gradients(group['params'])
-        return self.base_optimizer.step(closure)
+        return self.optimizer.step(closure)
diff --git a/pytorch_optimizer/optimizer/pcgrad.py b/pytorch_optimizer/optimizer/pcgrad.py
index f4b42889..1cd19723 100644
--- a/pytorch_optimizer/optimizer/pcgrad.py
+++ b/pytorch_optimizer/optimizer/pcgrad.py
@@ -5,9 +5,9 @@
 import numpy as np
 import torch
 from torch import nn
+from torch.optim import Optimizer
 
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import OPTIMIZER
 
 
 def flatten_grad(grads: List[torch.Tensor]) -> torch.Tensor:
@@ -29,11 +29,11 @@ def un_flatten_grad(grads: torch.Tensor, shapes: List[int]) -> List[torch.Tensor
 class PCGrad(BaseOptimizer):
     r"""Gradient Surgery for Multi-Task Learning.
 
-    :param optimizer: OPTIMIZER: optimizer instance.
+    :param optimizer: Optimizer: optimizer instance.
     :param reduction: str. reduction method.
     """
 
-    def __init__(self, optimizer: OPTIMIZER, reduction: str = 'mean'):
+    def __init__(self, optimizer: Optimizer, reduction: str = 'mean'):
         self.validate_options(reduction, 'reduction', ['mean', 'sum'])
 
         self.optimizer = optimizer
@@ -49,7 +49,7 @@ def zero_grad(self):
     def step(self):
         return self.optimizer.step()
 
-    def set_grad(self, grads: List[torch.Tensor]):
+    def set_grad(self, grads: List[torch.Tensor]) -> None:
         idx: int = 0
         for group in self.optimizer.param_groups:
             for p in group['params']:
@@ -121,7 +121,7 @@ def project_conflicting(self, grads: List[torch.Tensor], has_grads: List[torch.T
 
         return merged_grad
 
-    def pc_backward(self, objectives: Iterable[nn.Module]):
+    def pc_backward(self, objectives: Iterable[nn.Module]) -> None:
         r"""Calculate the gradient of the parameters.
 
         :param objectives: Iterable[nn.Module]. a list of objectives.
diff --git a/pytorch_optimizer/optimizer/sam.py b/pytorch_optimizer/optimizer/sam.py
index 771f7643..114002fe 100644
--- a/pytorch_optimizer/optimizer/sam.py
+++ b/pytorch_optimizer/optimizer/sam.py
@@ -6,6 +6,7 @@
 from torch.distributed import ReduceOp, all_reduce, get_world_size, is_initialized
 from torch.nn.parallel import DistributedDataParallel
 from torch.nn.utils import clip_grad_norm_
+from torch.optim import Optimizer
 
 from pytorch_optimizer.base.exception import NoClosureError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
@@ -54,7 +55,7 @@ def closure():
                 optimizer.zero_grad()
 
     :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
-    :param base_optimizer: Optimizer. base optimizer.
+    :param base_optimizer: OPTIMIZER. base optimizer.
     :param rho: float. size of the neighborhood for computing the max loss.
     :param adaptive: bool. element-wise Adaptive SAM.
     :param perturb_eps: float. eps for perturbation.
@@ -79,7 +80,7 @@ def __init__(
         defaults.update(kwargs)
         super().__init__(params, defaults)
 
-        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+        self.base_optimizer: Optimizer = base_optimizer(self.param_groups, **kwargs)
         self.param_groups = self.base_optimizer.param_groups
 
     def __str__(self) -> str:
@@ -187,7 +188,7 @@ def loss_fn(predictions, targets):
     def __init__(
         self,
         params: PARAMETERS,
-        base_optimizer: OPTIMIZER,
+        base_optimizer: Optimizer,
         model: nn.Module,
         rho_scheduler,
         alpha: float = 0.4,
diff --git a/pytorch_optimizer/optimizer/spam.py b/pytorch_optimizer/optimizer/spam.py
new file mode 100644
index 00000000..3221ac06
--- /dev/null
+++ b/pytorch_optimizer/optimizer/spam.py
@@ -0,0 +1,278 @@
+import math
+
+import torch
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class CosineDecay:
+    r"""Applies cosine decay to a parameter (death_rate), using PyTorch's built-in `CosineAnnealingLR`.
+
+    :param death_rate: float. initial value to be decayed.
+    :param t_max: int. maximum number of iterations for the decay.
+    :param eta_min: Optional[float]. minimum value of the parameter after decay. defaults to 0.
+    :param last_epoch: Optional[int]. the index of the last epoch. Defaults to -1.
+    """
+
+    def __init__(self, death_rate: float, t_max: int, eta_min: float = 0.0, last_epoch: int = -1):
+        self.sgd = torch.optim.SGD(
+            torch.nn.ParameterList([torch.nn.Parameter(torch.zeros(1))]),
+            lr=death_rate,
+        )
+        self.cosine_stepper = torch.optim.lr_scheduler.CosineAnnealingLR(self.sgd, t_max + 1, eta_min, last_epoch)
+        self.T_max = t_max
+        self.eta_min = eta_min
+
+    def step(self, current_step: int) -> None:
+        r"""One step of the cosine decay scheduler.
+
+        :param current_step: int. Current step index.
+        """
+        self.cosine_stepper.step(current_step)
+
+    def get_death_rate(self, current_step: int) -> float:
+        r"""Get the updated rate (death_rate) at the given step.
+
+        :param current_step: int. Current step index.
+        """
+        if current_step >= self.T_max:
+            return self.eta_min
+
+        self.step(current_step)
+
+        return self.sgd.param_groups[0]['lr']
+
+
+class SPAM(BaseOptimizer):
+    r"""Spike-Aware Adam with Momentum Reset for Stable LLM Training.
+
+    :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 density: float. density parameter. only used for 2d parameters (e.g. Linear).
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param warmup_epoch: int: number of epochs to warm up. defaults to 50.
+    :param threshold: int. threshold for gradient masking. defaults to 5000.
+    :param grad_accu_steps: int. gradient accumulation steps before threshold-based masking applies. defaults to 20.
+    :param update_proj_gap: int. update projection gap.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        betas: BETAS = (0.9, 0.999),
+        density: float = 1.0,
+        weight_decay: float = 0.0,
+        warmup_epoch: int = 150,
+        threshold: int = 5000,
+        grad_accu_steps: int = 20,
+        update_proj_gap: int = 500,
+        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(warmup_epoch, 'warmup_epoch')
+        self.validate_non_negative(density, 'density')
+        self.validate_non_negative(threshold, 'threshold')
+        self.validate_non_negative(grad_accu_steps, 'grad_accu_steps')
+        self.validate_non_negative(update_proj_gap, 'update_proj_gap')
+        self.validate_non_negative(eps, 'eps')
+
+        self.density = density
+        self.warmup_epoch = warmup_epoch
+        self.threshold = threshold
+        self.grad_accu_steps = grad_accu_steps
+        self.update_proj_gap = update_proj_gap
+        self.warmup = CosineDecay(0.99, warmup_epoch)
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'eps': eps,
+            **kwargs,
+        }
+        super().__init__(params, defaults)
+
+        self.init_masks()
+
+        self.state['total_step'] = 0
+        self.state['current_step'] = warmup_epoch + 1
+
+    @staticmethod
+    def initialize_random_rank_boolean_tensor(m: int, n: int, density: float) -> torch.Tensor:
+        r"""Create an (m x n) boolean tensor with `density` fraction of True entries.
+
+        :param m: int. number of rows.
+        :param n: int. number of columns.
+        :param density: float. fraction of True entries. 1.0 means all True.
+        """
+        total_elements: int = m * n
+        non_zero_count: int = int(density * total_elements)
+
+        tensor = torch.zeros((m, n), dtype=torch.bool)
+
+        if non_zero_count == 0:
+            return tensor
+
+        indices = torch.randperm(total_elements)[:non_zero_count]
+        rows, cols = indices // n, indices % n
+        tensor[rows, cols] = True
+
+        return tensor
+
+    def update_mask_random(self, density: float, p: torch.Tensor, old_mask: torch.Tensor) -> torch.Tensor:
+        r"""Update a random mask.
+
+        Create a new random mask with the same density, compute overlap ratio with old_mask, and update the EMA for
+        the overlap region.
+
+        :param density: float. fraction of elements to keep.
+        :param p: torch.Tensor. parameter to which the mask is applied.
+        :param old_mask: torch.Tensor. previous binary mask.
+        """
+        new_mask: torch.Tensor = torch.rand_like(p) < density
+
+        exp_avg = torch.zeros_like(p[new_mask])
+        exp_avg_sq = torch.zeros_like(p[new_mask])
+
+        intersection_mask = new_mask & old_mask
+        new_intersection_indices = intersection_mask[new_mask]
+        old_intersection_indices = intersection_mask[old_mask]
+
+        state = self.state[p]
+        exp_avg[new_intersection_indices] = state['exp_avg'][old_intersection_indices]
+        exp_avg_sq[new_intersection_indices] = state['exp_avg_sq'][old_intersection_indices]
+
+        state['exp_avg'] = exp_avg
+        state['exp_avg_sq'] = exp_avg_sq
+
+        return new_mask
+
+    def update_masks(self) -> None:
+        r"""Update masks in each parameter group that has 'density'.
+
+        The new mask is selected randomly, and the overlap ratio with the old mask is printed.
+        """
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+                if 'mask' in state:
+                    new_mask = self.update_mask_random(self.density, p, state['mask'])
+                    state['mask'] = new_mask
+                    p.mask = new_mask
+
+    def init_masks(self) -> None:
+        r"""Initialize random masks for each parameter group that has 'density'."""
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+                if p.dim() == 2 and 'mask' not in state:
+                    state['mask'] = self.initialize_random_rank_boolean_tensor(
+                        p.shape[0],
+                        p.shape[1],
+                        density=self.density,
+                    ).to(p.device)
+
+    def __str__(self) -> str:
+        return 'SPAM'
+
+    @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()
+
+        scale_factor: float = 1.0 - self.warmup.get_death_rate(self.state['current_step'])
+
+        for group in self.param_groups:
+            if 'step' not in group:
+                group['step'] = 1
+            else:
+                group['step'] += 1
+
+            beta1, beta2 = group['betas']
+
+            bias_correction1: float = self.debias(beta1, group['step'])
+            bias_correction2_sq: float = math.sqrt(self.debias(beta2, group['step']))
+
+            step_size: float = group['lr'] * 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 'mask' in state:
+                    grad = grad[state['mask']]
+
+                if 'exp_avg' not in state:
+                    state['exp_avg'] = torch.zeros_like(grad)
+                    state['exp_avg_sq'] = torch.zeros_like(grad)
+
+                if (self.state['total_step'] + 1) % self.update_proj_gap == 0:
+                    state['exp_avg'] = torch.zeros_like(grad)
+                    state['exp_avg_sq'] = torch.zeros_like(grad)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+
+                if self.threshold != 0:
+                    current_step: int = self.state['total_step'] + 1
+                    if current_step >= self.grad_accu_steps and (
+                        self.update_proj_gap == 0 or current_step % self.update_proj_gap >= self.grad_accu_steps
+                    ):
+                        mask = grad.pow(2) > (self.threshold * exp_avg_sq)
+                        grad[mask].sign_().mul_(torch.sqrt(exp_avg_sq[mask] * self.threshold))
+
+                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'])
+
+                if 'mask' in state:
+                    grad_full = torch.zeros_like(p.grad)
+                    grad_full[state['mask']] = exp_avg / de_nom
+                    p.add_(grad_full, alpha=-step_size * scale_factor)
+                else:
+                    p.addcdiv_(exp_avg, de_nom, value=-step_size * scale_factor)
+
+                self.apply_weight_decay(
+                    p[state['mask']] if 'mask' in state else p,
+                    grad=None,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=True,
+                    fixed_decay=False,
+                )
+
+        self.state['total_step'] += 1
+        self.state['current_step'] += 1
+
+        if (self.state['total_step'] != 0) and (self.state['total_step'] + 1) % self.update_proj_gap == 0:
+            self.update_masks()
+            self.state['current_step'] = 0
+            self.warmup = CosineDecay(0.99, self.warmup_epoch)
+
+        return loss
diff --git a/pytorch_optimizer/optimizer/trac.py b/pytorch_optimizer/optimizer/trac.py
index 8b51aa78..6982de36 100644
--- a/pytorch_optimizer/optimizer/trac.py
+++ b/pytorch_optimizer/optimizer/trac.py
@@ -2,9 +2,10 @@
 
 import torch
 from torch import nn
+from torch.optim import Optimizer
 
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, OPTIMIZER
+from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, OPTIMIZER_INSTANCE_OR_CLASS
 
 
 def polyval(x: torch.Tensor, coef: torch.Tensor) -> torch.Tensor:
@@ -91,7 +92,7 @@ class TRAC(BaseOptimizer):
 
                 optimizer.step()
 
-    :param optimizer: Optimizer. base optimizer.
+    :param optimizer: OPTIMIZER_INSTANCE_OR_CLASS. base optimizer.
     :param betas: List[float]. list of beta values.
     :param num_coefs: int. the number of polynomial coefficients to use in the approximation.
     :param s_prev: float. initial scale value.
@@ -100,7 +101,7 @@ class TRAC(BaseOptimizer):
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: OPTIMIZER_INSTANCE_OR_CLASS,
         betas: List[float] = (0.9, 0.99, 0.999, 0.9999, 0.99999, 0.999999),
         num_coefs: int = 128,
         s_prev: float = 1e-8,
@@ -111,6 +112,14 @@ def __init__(
         self.validate_non_negative(s_prev, 's_prev')
         self.validate_non_negative(eps, 'eps')
 
+        if isinstance(optimizer, Optimizer):
+            self.optimizer = optimizer
+        elif 'params' in kwargs:
+            params = kwargs.pop('params')
+            self.optimizer = optimizer(params, **kwargs)
+        else:
+            raise ValueError('Need to pass `params` when you pass the torch.optim.Optimizer instance.')
+
         self._optimizer_step_pre_hooks: Dict[int, Callable] = {}
         self._optimizer_step_post_hooks: Dict[int, Callable] = {}
 
@@ -121,8 +130,7 @@ def __init__(
 
         self.f_term = self.s_prev / self.erf_imag(1.0 / torch.sqrt(torch.tensor(2.0)))
 
-        self.optimizer = optimizer
-        self.defaults: DEFAULTS = optimizer.defaults
+        self.defaults: DEFAULTS = self.optimizer.defaults
 
     def __str__(self) -> str:
         return 'TRAC'


Test patch:
diff --git a/tests/constants.py b/tests/constants.py
index 6ed59335..32d195ac 100644
--- a/tests/constants.py
+++ b/tests/constants.py
@@ -18,6 +18,7 @@
     SGDW,
     SM3,
     SOAP,
+    SPAM,
     SRMM,
     SWATS,
     A2Grad,
@@ -547,8 +548,10 @@
     (SGDSaI, {'lr': 1e0}, 15),
     (SGDSaI, {'lr': 1e0, 'momentum': 0.0}, 15),
     (Grams, {'lr': 1e-1, 'weight_decay': 1e-3}, 5),
-    (Ranger25, {'lr': 1e-1}, 25),
-    (Ranger25, {'lr': 1e-1, 't_alpha_beta3': 5}, 25),
+    (SPAM, {'lr': 1e0, 'weight_decay': 1e-3, 'warmup_epoch': 1, 'grad_accu_steps': 1, 'update_proj_gap': 1}, 5),
+    (Ranger25, {'lr': 1e0}, 5),
+    (Ranger25, {'lr': 1e0, 't_alpha_beta3': 5}, 5),
+    (Ranger25, {'lr': 1e-1, 'stable_adamw': False, 'eps': None}, 5),
 ]
 ADANORM_SUPPORTED_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'adanorm': True}, 10),
diff --git a/tests/test_create_optimizer.py b/tests/test_create_optimizer.py
index 09228ae0..e13172e8 100644
--- a/tests/test_create_optimizer.py
+++ b/tests/test_create_optimizer.py
@@ -6,11 +6,15 @@
 
 
 @pytest.mark.parametrize('use_lookahead', [True, False])
+@pytest.mark.parametrize('use_orthograd', [True, False])
 @pytest.mark.parametrize('optimizer_name', VALID_OPTIMIZER_NAMES)
-def test_create_optimizer(use_lookahead, optimizer_name):
+def test_create_optimizer(use_lookahead, use_orthograd, optimizer_name):
     if optimizer_name in ('adamw', 'demo'):
         pytest.skip(f'skip {optimizer_name}')
 
+    if use_lookahead and use_orthograd:
+        pytest.skip()
+
     kwargs = {'eps': 1e-8, 'k': 7}
     if optimizer_name == 'ranger21':
         kwargs.update({'num_iterations': 1})
@@ -23,6 +27,7 @@ def test_create_optimizer(use_lookahead, optimizer_name):
         LogisticRegression(),
         optimizer_name=optimizer_name,
         use_lookahead=use_lookahead,
+        use_orthograd=use_orthograd,
         **kwargs,
     )
 
diff --git a/tests/test_gradients.py b/tests/test_gradients.py
index bdf24741..a8e27372 100644
--- a/tests/test_gradients.py
+++ b/tests/test_gradients.py
@@ -2,12 +2,12 @@
 import torch
 
 from pytorch_optimizer.base.exception import NoSparseGradientError
-from pytorch_optimizer.optimizer import SAM, TRAC, WSAM, AdamP, Lookahead, load_optimizer
+from pytorch_optimizer.optimizer import SAM, TRAC, WSAM, AdamP, Lookahead, OrthoGrad, load_optimizer
 from tests.constants import NO_SPARSE_OPTIMIZERS, SPARSE_OPTIMIZERS, VALID_OPTIMIZER_NAMES
 from tests.utils import build_environment, simple_parameter, simple_sparse_parameter, sphere_loss
 
 
-@pytest.mark.parametrize('optimizer_name', [*VALID_OPTIMIZER_NAMES, 'lookahead', 'trac'])
+@pytest.mark.parametrize('optimizer_name', [*VALID_OPTIMIZER_NAMES, 'lookahead', 'trac', 'orthograd'])
 def test_no_gradients(optimizer_name):
     if optimizer_name in {'lomo', 'adalomo', 'adammini', 'demo'}:
         pytest.skip(f'skip {optimizer_name} optimizer.')
@@ -28,6 +28,8 @@ def test_no_gradients(optimizer_name):
         optimizer = Lookahead(load_optimizer('adamw')(params), k=1)
     elif optimizer_name == 'trac':
         optimizer = TRAC(load_optimizer('adamw')(params))
+    elif optimizer_name == 'orthograd':
+        optimizer = OrthoGrad(load_optimizer('adamw')(params))
     else:
         optimizer = load_optimizer(optimizer_name)(params)
 
diff --git a/tests/test_load_modules.py b/tests/test_load_modules.py
index 42b43c8e..df2ed9e0 100644
--- a/tests/test_load_modules.py
+++ b/tests/test_load_modules.py
@@ -34,7 +34,7 @@ def test_load_lr_scheduler_invalid(invalid_lr_scheduler_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 87
+    assert len(get_supported_optimizers()) == 88
     assert len(get_supported_optimizers('adam*')) == 7
     assert len(get_supported_optimizers(['adam*', 'ranger*'])) == 10
 
diff --git a/tests/test_optimizer_parameters.py b/tests/test_optimizer_parameters.py
index 1e2d6259..3a001cf7 100644
--- a/tests/test_optimizer_parameters.py
+++ b/tests/test_optimizer_parameters.py
@@ -2,7 +2,17 @@
 import torch
 from torch import nn
 
-from pytorch_optimizer.optimizer import SAM, WSAM, Lookahead, PCGrad, Ranger21, SafeFP16Optimizer, load_optimizer
+from pytorch_optimizer.optimizer import (
+    SAM,
+    TRAC,
+    WSAM,
+    Lookahead,
+    OrthoGrad,
+    PCGrad,
+    Ranger21,
+    SafeFP16Optimizer,
+    load_optimizer,
+)
 from pytorch_optimizer.optimizer.galore_utils import GaLoreProjector
 from tests.constants import PULLBACK_MOMENTUM
 from tests.utils import Example, simple_parameter, simple_zero_rank_parameter
@@ -44,7 +54,7 @@ def test_came_parameters():
 
 
 def test_pcgrad_parameters():
-    opt = load_optimizer('adamw')([simple_parameter()])
+    opt = load_optimizer('adamp')([simple_parameter()])
 
     # test reduction
     for reduction in ['mean', 'sum']:
@@ -56,16 +66,17 @@ def test_pcgrad_parameters():
 
 def test_sam_parameters():
     with pytest.raises(ValueError):
-        SAM(None, load_optimizer('adamw'), rho=-0.1)
+        SAM(None, load_optimizer('adamp'), rho=-0.1)
 
 
 def test_wsam_parameters():
     with pytest.raises(ValueError):
-        WSAM(None, None, load_optimizer('adamw'), rho=-0.1)
+        WSAM(None, None, load_optimizer('adamp'), rho=-0.1)
 
 
 def test_lookahead_parameters():
-    optimizer = load_optimizer('adamw')([simple_parameter()])
+    optimizer_instance = load_optimizer('adamp')
+    optimizer = optimizer_instance([simple_parameter()])
 
     for pullback_momentum in PULLBACK_MOMENTUM:
         opt = Lookahead(optimizer, pullback_momentum=pullback_momentum)
@@ -267,3 +278,14 @@ def test_galore_projection_type():
 
     with pytest.raises(ValueError):
         GaLoreProjector.get_orthogonal_matrix(p, 1, projection_type='std')
+
+
+@pytest.mark.parametrize('optimizer_instance', [Lookahead, OrthoGrad, TRAC])
+def test_load_optimizer(optimizer_instance):
+    params = [simple_parameter()]
+
+    _ = optimizer_instance(torch.optim.AdamW(params))
+    _ = optimizer_instance(torch.optim.AdamW, params=params)
+
+    with pytest.raises(ValueError):
+        optimizer_instance(torch.optim.AdamW)
diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py
index d9d39550..50dd73bf 100644
--- a/tests/test_optimizers.py
+++ b/tests/test_optimizers.py
@@ -64,7 +64,6 @@ def _closure() -> float:
     (x_data, y_data), model, loss_fn = environment
 
     optimizer_class, config, iterations = optimizer_fp32_config
-
     optimizer_name: str = optimizer_class.__name__
     if optimizer_name == 'Nero' and 'constraints' not in config:
         pytest.skip(f'skip {optimizer_name} w/o {config}')
@@ -846,3 +845,11 @@ def test_mars_c_t_norm():
 
     optimizer = load_optimizer('mars')([param], optimize_1d=True)
     optimizer.step()
+
+
+def test_spam_optimizer():
+    optimizer = load_optimizer('spam')(Example().parameters(), density=0.0)
+    optimizer.step()
+
+    optimizer = load_optimizer('spam')([simple_parameter(True)], grad_accu_steps=0, update_proj_gap=1)
+    optimizer.step()
diff --git a/tests/test_utils.py b/tests/test_utils.py
index b914731d..462c84ee 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -265,4 +265,6 @@ def test_cpu_offload_optimizer():
 
 def test_orthograd_name():
     optimizer = build_orthograd(Example().parameters())
+    _ = optimizer.param_groups
+    _ = optimizer.__getstate__()
     assert str(optimizer).lower() == 'orthograd'
diff --git a/tests/utils.py b/tests/utils.py
index 71528a22..1deeebd8 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -108,7 +108,7 @@ def build_lookahead(*parameters, **kwargs):
 
 
 def build_orthograd(*parameters, **kwargs):
-    return OrthoGrad(*parameters, torch.optim.AdamW, **kwargs)
+    return OrthoGrad(AdamW(*parameters, **kwargs))
 
 
 def ids(v) -> str: