pytorch-optimizer 3.5.0__tar.gz → 3.5.1__tar.gz

{pytorch_optimizer-3.5.0 → pytorch_optimizer-3.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pytorch_optimizer
-Version: 3.5.0
+Version: 3.5.1
 Summary: optimizer & lr scheduler & objective function collections in PyTorch
 License: Apache-2.0
 Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaGC,AdaMax,AdamG,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdEMAMix,Simplified-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,EXAdam,FAdam,FOCUS,Fromage,FTRL,GaLore,Grams,Gravity,GrokFast,GSAM,Kate,Lamb,LaProp,LARS,Lion,LOMO,Lookahead,MADGRAD,MARS,MSVAG,Muno,Nero,NovoGrad,OrthoGrad,PAdam,PCGrad,PID,PNM,Prodigy,PSGD,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,GCSAM,LookSAM,ScheduleFreeSGD,ScheduleFreeAdamW,ScheduleFreeRAdam,SCION,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SOAP,SopihaH,SPAM,StableSPAM,SRMM,StableAdamW,SWATS,TAM,Tiger,TRAC,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD,QGaLore
@@ -244,7 +244,7 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | EXAdam              | *The Power of Adaptive Cross-Moments*                                                             | [github](https://github.com/AhmedMostafa16/EXAdam)                                                             | <https://arxiv.org/abs/2412.20302>                                                          | [cite](https://github.com/AhmedMostafa16/EXAdam?tab=readme-ov-file#citation)                                                        |
 | GCSAM               | *Gradient Centralized Sharpness Aware Minimization*                                               | [github](https://github.com/mhassann22/GCSAM)                                                                  | <https://arxiv.org/abs/2501.11584>                                                          | [cite](https://github.com/mhassann22/GCSAM?tab=readme-ov-file#citation)                                                             |
 | LookSAM             | *Towards Efficient and Scalable Sharpness-Aware Minimization*                                     | [github](https://github.com/rollovd/LookSAM)                                                                   | <https://arxiv.org/abs/2203.02714>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2022arXiv220302714L/exportcitation)                                                        |
-| SCION               | *Training Deep Learning Models with Norm-Constrained LMOs*                                        |                                                                                                                | <https://arxiv.org/abs/2502.07529>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250207529P/exportcitation)                                                        |
+| SCION               | *Training Deep Learning Models with Norm-Constrained LMOs*                                        | [github](https://github.com/LIONS-EPFL/scion)                                                                  | <https://arxiv.org/abs/2502.07529>                                                          | [cite](https://github.com/LIONS-EPFL/scion?tab=readme-ov-file#citation)                                                             |
 | COSMOS              | *SOAP with Muon*                                                                                  | [github](https://github.com/lliu606/COSMOS)                                                                    |                                                                                             |                                                                                                                                     |
 | StableSPAM          | *How to Train in 4-Bit More Stably than 16-Bit Adam*                                              | [github](https://github.com/TianjinYellow/StableSPAM)                                                          | <https://arxiv.org/abs/2502.17055>                                                          |                                                                                                                                     |
 | AdaGC               | *Improving Training Stability for Large Language Model Pretraining*                               |                                                                                                                | <https://arxiv.org/abs/2502.11034>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250211034W/exportcitation)                                                        |

{pytorch_optimizer-3.5.0 → pytorch_optimizer-3.5.1}/README.md

@@ -206,7 +206,7 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | EXAdam              | *The Power of Adaptive Cross-Moments*                                                             | [github](https://github.com/AhmedMostafa16/EXAdam)                                                             | <https://arxiv.org/abs/2412.20302>                                                          | [cite](https://github.com/AhmedMostafa16/EXAdam?tab=readme-ov-file#citation)                                                        |
 | GCSAM               | *Gradient Centralized Sharpness Aware Minimization*                                               | [github](https://github.com/mhassann22/GCSAM)                                                                  | <https://arxiv.org/abs/2501.11584>                                                          | [cite](https://github.com/mhassann22/GCSAM?tab=readme-ov-file#citation)                                                             |
 | LookSAM             | *Towards Efficient and Scalable Sharpness-Aware Minimization*                                     | [github](https://github.com/rollovd/LookSAM)                                                                   | <https://arxiv.org/abs/2203.02714>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2022arXiv220302714L/exportcitation)                                                        |
-| SCION               | *Training Deep Learning Models with Norm-Constrained LMOs*                                        |                                                                                                                | <https://arxiv.org/abs/2502.07529>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250207529P/exportcitation)                                                        |
+| SCION               | *Training Deep Learning Models with Norm-Constrained LMOs*                                        | [github](https://github.com/LIONS-EPFL/scion)                                                                  | <https://arxiv.org/abs/2502.07529>                                                          | [cite](https://github.com/LIONS-EPFL/scion?tab=readme-ov-file#citation)                                                             |
 | COSMOS              | *SOAP with Muon*                                                                                  | [github](https://github.com/lliu606/COSMOS)                                                                    |                                                                                             |                                                                                                                                     |
 | StableSPAM          | *How to Train in 4-Bit More Stably than 16-Bit Adam*                                              | [github](https://github.com/TianjinYellow/StableSPAM)                                                          | <https://arxiv.org/abs/2502.17055>                                                          |                                                                                                                                     |
 | AdaGC               | *Improving Training Stability for Large Language Model Pretraining*                               |                                                                                                                | <https://arxiv.org/abs/2502.11034>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250211034W/exportcitation)                                                        |

{pytorch_optimizer-3.5.0 → pytorch_optimizer-3.5.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.5.0"
+version = "3.5.1"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]

{pytorch_optimizer-3.5.0 → pytorch_optimizer-3.5.1}/pytorch_optimizer/__init__.py

@@ -141,6 +141,7 @@ from pytorch_optimizer.optimizer import (
     ScheduleFreeRAdam,
     ScheduleFreeSGD,
     ScheduleFreeWrapper,
+    SCIONLight,
     SGDSaI,
     Shampoo,
     SignSGD,

{pytorch_optimizer-3.5.0 → pytorch_optimizer-3.5.1}/pytorch_optimizer/optimizer/__init__.py

@@ -87,7 +87,7 @@ from pytorch_optimizer.optimizer.schedulefree import (
     ScheduleFreeSGD,
     ScheduleFreeWrapper,
 )
-from pytorch_optimizer.optimizer.scion import SCION
+from pytorch_optimizer.optimizer.scion import SCION, SCIONLight
 from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD, SGDSaI, SignSGD
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo
@@ -309,6 +309,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     Kron,
     EXAdam,
     SCION,
+    SCIONLight,
     StableSPAM,
     AdaGC,
     Ranger25,

{pytorch_optimizer-3.5.0 → pytorch_optimizer-3.5.1}/pytorch_optimizer/optimizer/muon.py

@@ -1,6 +1,6 @@
 import math
 import os
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import torch
 from torch.distributed import ReduceOp, all_reduce
@@ -131,9 +131,18 @@ class Muon(BaseOptimizer):
                 state['moment2'] = torch.zeros_like(p)
 
     @staticmethod
-    def adjust_lr_for_muon(lr: float, param_shape) -> float:
-        adjusted_ratio: float = 0.2 * math.sqrt(max(param_shape[0], param_shape[1]))
-        return lr * adjusted_ratio
+    def get_adjusted_lr(lr: float, param_shape: Tuple[float, ...], use_adjusted_lr: bool = False) -> float:
+        r"""Get the adjust learning rate."""
+        output_shape, *input_shape = param_shape
+        input_shape = math.prod(input_shape)
+
+        ratio: float = (
+            math.pow(max(1.0, output_shape / input_shape), 0.5)
+            if use_adjusted_lr
+            else 0.2 * math.sqrt(max(output_shape, input_shape))
+        )
+
+        return lr * ratio
 
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
@@ -202,9 +211,9 @@ class Muon(BaseOptimizer):
                     fixed_decay=False,
                 )
 
-                lr: float = self.adjust_lr_for_muon(group['lr'], p.size()) if group['use_adjusted_lr'] else group['lr']
+                lr: float = self.get_adjusted_lr(group['lr'], p.size(), group['use_adjusted_lr'])
 
-                p.add_(g, alpha=-lr * (max(1.0, p.size(-2) / p.size(-1)) ** 0.5))
+                p.add_(g, alpha=-lr)
                 curr_idx += p.numel()
 
             params = [p for p in group['params'] if p.grad is not None and not self.state[p]['use_muon']]

pytorch_optimizer-3.5.1/pytorch_optimizer/optimizer/scion.py

@@ -0,0 +1,530 @@
+import math
+from enum import IntEnum
+from typing import Dict, Optional
+
+import torch
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.type import CLOSURE, DEFAULTS, LOSS, PARAMETERS
+from pytorch_optimizer.optimizer.shampoo_utils import zero_power_via_newton_schulz_5
+
+
+class LMONorm(IntEnum):
+    r"""normalization types."""
+
+    NONE = 0
+    AUTO = 1
+    SPECTRAL = 2
+    SPECTRALCONV = 3
+    SIGN = 4
+    BIAS = 5
+    COL = 6
+    ROW = 7
+
+
+class Norm:
+    r"""Base class to perform norm onto Scion. This class does no norm."""
+
+    def init(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Initialize parameter."""
+        return x
+
+    def lmo(self, grad: torch.Tensor) -> torch.Tensor:
+        r"""Get LMO."""
+        return grad
+
+
+class Col(Norm):
+    r"""col-wise normalization.
+
+    :param normalized: bool. normalize by the input dimension. use for non-input layers.
+    :param transpose: bool. transpose input before normalization. use for embedding layers which have a shape of
+        (vocab_size, embedding_dim)
+    """
+
+    def __init__(self, normalized: bool = False, transpose: bool = False) -> None:
+        self.normalized = normalized
+        self.transpose = transpose
+
+    def init(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        if self.transpose:
+            x = x.transpose(0, 1)
+
+        torch.nn.init.normal_(x)
+
+        x.div_(x.norm(dim=0, keepdim=True)).mul_(math.sqrt(x.size(0)))
+        if self.normalized:
+            x.div_(x.size(1))
+
+        x = x.to(dtype=dtype)
+        if self.transpose:
+            x = x.transpose(0, 1)
+
+        return x
+
+    def lmo(self, grad: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+        if self.transpose:
+            grad = grad.transpose(0, 1)
+
+        d_in, d_out = grad.size()
+
+        rms_value = torch.sqrt(torch.sum(grad.pow(2), dim=0, keepdim=True)) / math.sqrt(d_in)
+        if self.normalized:
+            rms_value.mul_(d_out)
+
+        grad /= rms_value.add_(eps)
+
+        if self.transpose:
+            grad = grad.transpose(0, 1)
+
+        return grad
+
+
+class Row(Norm):
+    r"""row-wise normalization.
+
+    :param normalized: bool. normalize by the input dimension. use for non-input layers.
+    :param transpose: bool. transpose input before normalization. use for embedding layers which have a shape of
+        (vocab_size, embedding_dim)
+    """
+
+    def __init__(self, normalized: bool = True, transpose: bool = False) -> None:
+        self.normalized = normalized
+        self.transpose = transpose
+
+    def init(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        if self.transpose:
+            x = x.transpose(0, 1)
+
+        torch.nn.init.normal_(x)
+
+        x.div_(x.norm(dim=-1, keepdim=True))
+        if self.normalized:
+            x.div_(math.sqrt(x.size(-1)))
+
+        x = x.to(dtype=dtype)
+        if self.transpose:
+            x = x.transpose(0, 1)
+
+        return x
+
+    def lmo(self, grad: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+        if self.transpose:
+            grad = grad.transpose(0, 1)
+
+        rms_value = torch.sqrt(torch.sum(grad.pow(2), dim=-1, keepdim=True))
+        if self.normalized:
+            rms_value.mul_(math.sqrt(grad.size(-1)))
+
+        grad /= rms_value.add_(eps)
+
+        if self.transpose:
+            grad = grad.transpose(0, 1)
+
+        return grad
+
+
+class BiasRMS(Norm):
+    r"""bias RMS."""
+
+    def init(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.init.zeros_(x)
+
+    def lmo(self, grad: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+        rms_value = torch.sqrt(torch.sum(grad.pow(2), dim=0, keepdim=True))
+        grad /= rms_value.add_(eps)
+        return grad
+
+
+class SpectralConv(Norm):
+    r"""spectral-convolution normalization.
+
+    :param num_steps: int. number of steps of zero-power Newton-Schulz 5.
+    """
+
+    def __init__(self, num_steps: int = 5) -> None:
+        self.num_steps = num_steps
+
+    def init(self, x: torch.Tensor) -> torch.Tensor:
+        x_fp64 = x.double()
+
+        d_out, d_in, kernel_size, *_ = x_fp64.size()
+
+        for i in range(kernel_size):
+            for j in range(kernel_size):
+                torch.nn.init.orthogonal_(x_fp64[..., i, j])
+
+        x_fp64.mul_(math.sqrt(d_out / d_in) / (kernel_size**2))
+
+        return x_fp64.to(dtype=x.dtype)
+
+    def lmo(self, grad: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+        grad = zero_power_via_newton_schulz_5(grad.view(len(grad), -1), self.num_steps).view(grad.shape)
+
+        d_out, d_in, kernel_size, *_ = grad.size()
+
+        grad *= math.sqrt(d_out / d_in) / (kernel_size**2)
+
+        return grad
+
+
+class Spectral(Norm):
+    r"""spectral normalization.
+
+    :param max_scale: bool. set upper bound (1.0) of the scale.
+    :param normalize: bool. normalize by the input dimension. use for non-input layers.
+    :param num_steps: int. number of steps of zero-power Newton-Schulz 5.
+    """
+
+    def __init__(self, max_scale: bool = False, normalize: bool = True, num_steps: int = 5) -> None:
+        self.max_scale = max_scale
+        self.normalize = normalize
+        self.num_steps = num_steps
+
+    def init(self, x: torch.Tensor) -> torch.Tensor:
+        x_fp64 = x.double()
+
+        torch.nn.init.orthogonal_(x_fp64)
+
+        d_out, d_in = x_fp64.size()
+
+        scale: float = math.sqrt(d_out / d_in) if self.normalize else math.sqrt(d_out)
+        if self.max_scale:
+            scale = max(1.0, scale)
+
+        x_fp64.mul_(scale)
+
+        return x_fp64.to(dtype=x.dtype)
+
+    def lmo(self, grad: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+        grad = zero_power_via_newton_schulz_5(grad.view(len(grad), -1), self.num_steps).view(grad.shape)
+
+        d_out, d_in = grad.size()
+
+        scale: float = math.sqrt(d_out / d_in) if self.normalize else math.sqrt(d_out)
+        if self.max_scale:
+            scale = max(1.0, scale)
+
+        grad *= scale
+
+        return grad
+
+
+class Sign(Norm):
+    r"""sign normalization.
+
+    :param zero_init: bool. initialize with zero.
+    :param normalize: bool. normalize by the input dimension. use for non-input layers.
+    """
+
+    def __init__(self, zero_init: bool = False, normalize: bool = True) -> None:
+        self.zero_init = zero_init
+        self.normalize = normalize
+
+    def init(self, x: torch.Tensor) -> torch.Tensor:
+        if self.zero_init:
+            return torch.nn.init.zeros_(x)
+
+        d_in: int = x.size(1)
+
+        x = 2 * torch.randint(0, 2, x.shape, dtype=x.dtype, device=x.device) - 1
+        if self.normalize:
+            x.div_(d_in)
+
+        return x
+
+    def lmo(self, grad: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+        d_in: int = grad.size(1)
+        return torch.sign(grad).div_(d_in) if self.normalize else torch.sign(grad)
+
+
+class Auto(Norm):
+    r"""choose Norm type automatically."""
+
+    def init(self, x: torch.Tensor) -> torch.Tensor:
+        ndim: int = x.ndim
+        if ndim in (0, 1):
+            return BiasRMS().init(x)
+        if ndim == 2:
+            return Spectral().init(x)
+        if ndim in (3, 4):
+            return SpectralConv().init(x)
+        raise NotImplementedError
+
+    def lmo(self, grad: torch.Tensor) -> torch.Tensor:
+        ndim: int = grad.ndim
+        if ndim in (0, 1):
+            return BiasRMS().lmo(grad)
+        if ndim == 2:
+            return Spectral().lmo(grad)
+        if ndim in (3, 4):
+            return SpectralConv().lmo(grad)
+        raise NotImplementedError
+
+
+def build_lmo_norm(norm_type: int, **kwargs) -> Norm:  # noqa: PLR0911
+    r"""Build LMONorm by given norm_type."""
+    if norm_type == LMONorm.AUTO:
+        return Auto()
+    if norm_type == LMONorm.SPECTRAL:
+        return Spectral(**kwargs)
+    if norm_type == LMONorm.SPECTRALCONV:
+        return SpectralConv(**kwargs)
+    if norm_type == LMONorm.SIGN:
+        return Sign(**kwargs)
+    if norm_type == LMONorm.BIAS:
+        return BiasRMS()
+    if norm_type == LMONorm.COL:
+        return Col(**kwargs)
+    if norm_type == LMONorm.ROW:
+        return Row(**kwargs)
+    return Norm()
+
+
+class SCION(BaseOptimizer):
+    r"""Training Deep Learning Models with Norm-Constrained LMOs.
+
+    Example:
+        >>> radius = 50.0
+        >>> parameter_groups = [{
+        ...     'params': model.transformer.h.parameters(),
+        ...     'norm_type': 'spectral',
+        ...     'norm_kwargs': {},
+        ...     'scale': radius,
+        ... }, {
+        ...     'params': model.lm_head.parameters(),
+        ...     'norm_type': 'sign',
+        ...     'norm_kwargs': {},
+        ...     'scale': radius * 60.0,
+        ... }]
+        >>> optimizer = SCION(parameter_groups)
+
+        For more details, checkout here https://github.com/LIONS-EPFL/scion/tree/main?tab=readme-ov-file#examples
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param momentum: float. momentum factor. 1.0 - usual momentum.
+    :param constraint: bool. whether to use a constraint SCG or not.
+    :param norm_type: int. supported LMO norm types. 0 stands for no normalization and 1 stands for AUTO. 0 to 7.
+        please check LMONorm Enum class for the details.
+    :param norm_kwargs: Optional[Dict]. arguments for the Norm.
+    :param scale: float. based on the usage of the original intend, 50.0 is used for Transformer block, and 3000.0 is
+        used for others (e.g. Embedding, LM head)
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        momentum: float = 0.1,
+        constraint: bool = False,
+        norm_type: int = LMONorm.AUTO,
+        norm_kwargs: Optional[Dict] = None,
+        scale: float = 1.0,
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_range(momentum, 'momentum', 0.0, 1.0, '(]')
+        self.validate_positive(scale, 'scale')
+
+        if norm_kwargs is None:
+            norm_kwargs = {}
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'momentum': momentum,
+            'constraint': constraint,
+            'norm_type': norm_type,
+            'norm_kwargs': norm_kwargs,
+            'scale': scale,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'SCION'
+
+    @torch.no_grad()
+    def reset(self):
+        pass
+
+    @torch.no_grad()
+    def init(self):
+        for group in self.param_groups:
+            norm = build_lmo_norm(group['norm_type'], **group['norm_kwargs'])
+            for p in group['params']:
+                norm.init(p)
+                p.mul_(group['scale'])
+
+    @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:
+            norm = build_lmo_norm(group['norm_type'], **group['norm_kwargs'])
+
+            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 'd' not in state:
+                    state['d'] = torch.zeros_like(grad)
+
+                d = state['d']
+                d.mul_(1.0 - group['momentum']).add_(grad, alpha=group['momentum'])
+
+                update = norm.lmo(grad).mul_(group['scale'])
+
+                if group['constraint']:
+                    p.mul_(1.0 - group['lr'])
+
+                if not group['constraint'] and group['weight_decay'] > 0.0:
+                    self.apply_weight_decay(
+                        p,
+                        grad,
+                        lr=group['lr'],
+                        weight_decay=group['weight_decay'],
+                        weight_decouple=group['weight_decouple'],
+                        fixed_decay=False,
+                    )
+
+                p.add_(update, alpha=-group['lr'])
+
+        return loss
+
+
+class SCIONLight(BaseOptimizer):
+    r"""Memory-efficient variant of the Scion optimizer.
+
+    Example:
+        >>> radius = 50.0
+        >>> parameter_groups = [{
+        ...     'params': model.transformer.h.parameters(),
+        ...     'norm_type': 'spectral',
+        ...     'norm_kwargs': {},
+        ...     'scale': radius,
+        ... }, {
+        ...     'params': model.lm_head.parameters(),
+        ...     'norm_type': 'sign',
+        ...     'norm_kwargs': {},
+        ...     'scale': radius * 60.0,
+        ... }]
+        >>> optimizer = SCIONLight(parameter_groups)
+
+        For more details, checkout here https://github.com/LIONS-EPFL/scion/tree/main?tab=readme-ov-file#examples
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param momentum: float. momentum factor. 1.0 - usual momentum.
+    :param constraint: bool. whether to use a constraint SCG or not.
+    :param norm_type: int. supported LMO norm types. 0 stands for no normalization and 1 stands for AUTO. 0 to 7.
+        please check LMONorm Enum class for the details.
+    :param norm_kwargs: Optional[Dict]. arguments for the Norm.
+    :param scale: float. based on the usage of the original intend, 50.0 is used for Transformer block, and 3000.0 is
+        used for others (e.g. Embedding, LM head)
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        momentum: float = 0.1,
+        constraint: bool = False,
+        norm_type: int = LMONorm.AUTO,
+        norm_kwargs: Optional[Dict] = None,
+        scale: float = 1.0,
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_range(momentum, 'momentum', 0.0, 1.0, '(]')
+        self.validate_positive(scale, 'scale')
+
+        if norm_kwargs is None:
+            norm_kwargs = {}
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'momentum': momentum,
+            'constraint': constraint,
+            'norm_type': norm_type,
+            'norm_kwargs': norm_kwargs,
+            'scale': scale,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'SCIONLight'
+
+    @torch.no_grad()
+    def reset(self):
+        pass
+
+    @torch.no_grad()
+    def init(self):
+        for group in self.param_groups:
+            norm = build_lmo_norm(group['norm_type'], **group['norm_kwargs'])
+            for p in group['params']:
+                norm.init(p)
+                p.mul_(group['scale'])
+
+    @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:
+            norm = build_lmo_norm(group['norm_type'], **group['norm_kwargs'])
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(str(self))
+
+                update = norm.lmo(grad).mul_(group['scale'])
+
+                if group['constraint']:
+                    p.mul_(1.0 - group['lr'])
+
+                if not group['constraint'] and group['weight_decay'] > 0.0:
+                    self.apply_weight_decay(
+                        p,
+                        grad,
+                        lr=group['lr'],
+                        weight_decay=group['weight_decay'],
+                        weight_decouple=group['weight_decouple'],
+                        fixed_decay=False,
+                    )
+
+                p.add_(update, alpha=-group['lr'])
+
+                if group['momentum'] != 1.0:
+                    grad.mul_(1.0 - group['momentum'])
+
+        return loss

{pytorch_optimizer-3.5.0 → pytorch_optimizer-3.5.1}/pytorch_optimizer/optimizer/spam.py

@@ -31,7 +31,8 @@ class CosineDecay:
 
         :param current_step: int. Current step index.
         """
-        self.cosine_stepper.step(current_step)
+        self.cosine_stepper.last_epoch = current_step
+        self.cosine_stepper.step()
 
     def get_death_rate(self, current_step: int) -> float:
         r"""Get the updated rate (death_rate) at the given step.
@@ -266,9 +267,9 @@ class StableSPAM(BaseOptimizer):
     :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 gamma1: float.
-    :param gamma2: float.
-    :param theta: float.
+    :param gamma1: float. gamma1 parameter.
+    :param gamma2: float. gamma2 parameter.
+    :param theta: float. theta parameter.
     :param t_max: Optional[int]. total number of steps.
     :param eta_min: float. eta_min of CosineDecay.
     :param weight_decay: float. weight decay (L2 penalty).