adv-optm 1.1.2__py3-none-any.whl → 1.2.dev1__py3-none-any.whl

adv_optm/__init__.py

@@ -5,6 +5,7 @@ from .optim import (
     Simplified_AdEMAMix,
     Lion_adv,
     Lion_Prodigy_adv,
+    Muon_adv,
 )
 
 __all__ = [
@@ -14,6 +15,7 @@ __all__ = [
     "Simplified_AdEMAMix",
     "Lion_adv",
     "Lion_Prodigy_adv",
+    "Muon_adv",
 ]
 
-__version__ = "1.1.2"
+__version__ = "1.2.dev1"

adv_optm/optim/AdamW_adv.py

@@ -209,7 +209,7 @@ class AdamW_adv(torch.optim.Optimizer):
         beta1, beta2 = group['betas']
 
         current_step = state['step']
-        if group['kourkoutas_beta']:
+        if group.get('kourkoutas_beta', False):
             # Call prepare_step() once at the beginning of the step for all params
             self.kourkoutas_helper.maybe_prepare_step(current_step)
             # Accumulate current grad's norm for the *next* step
@@ -220,7 +220,7 @@ class AdamW_adv(torch.optim.Optimizer):
         step = state['step'] + 1
         if group['use_bias_correction']:
             bias_correction1 = 1.0 - beta1 ** step
-            if group['kourkoutas_beta']:
+            if group.get('kourkoutas_beta', False):
                 bias_correction2 = 1.0 - group['betas'][1] ** step
                 # Use beta2_max for bias correction
             else:

adv_optm/optim/Adopt_adv.py

@@ -240,7 +240,7 @@ class Adopt_adv(torch.optim.Optimizer):
         beta1, beta2 = group['betas']
 
         current_step = state['step']
-        if group['kourkoutas_beta']:
+        if group.get('kourkoutas_beta', False):
             # Call prepare_step() once at the beginning of the step for all params
             self.kourkoutas_helper.maybe_prepare_step(current_step)
             # Accumulate current grad's norm for the *next* step

adv_optm/optim/Lion_Prodigy_adv.py

@@ -325,7 +325,7 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
             d_hat = self.d
             if global_d_denom > 0:
                 d_hat = d_coef * global_d_numerator / global_d_denom
-                if g_group['d_limiter']:
+                if g_group.get('d_limiter', False):
                     d_hat = min(self.d * (2 ** 0.25), d_hat)
                 if self.d == g_group['d0']:
                     self.d = max(self.d, d_hat)

adv_optm/optim/Muon_adv.py

@@ -0,0 +1,247 @@
+import torch
+from typing import Optional
+
+from ..util.BF16_Stochastic_Rounding import add_stochastic_
+from ..util.Newton_Schulz import _newton_schulz_iteration
+from ..util.Effective_Shape import _get_effective_shape
+from ..util.NNMF import _nnmf,_unnmf
+from ..util.One_Bit_Boolean import _pack_bools, _unpack_bools
+
+class Muon_adv(torch.optim.Optimizer):
+    """
+    Implements an advanced Muon algorithm.
+
+    Muon (MomentUm Orthogonalized by Newton-Schulz) is an optimizer designed for
+    the hidden layers of neural networks. It applies SGD with momentum and then
+    orthogonalizes the resulting update matrix using a Newton-Schulz iteration.
+
+    This implementation is designed for 2D parameters (e.g., linear layers) and
+    can handle other-dimensional parameters (e.g., 1D bias, 4D convolutional layers) by
+    flattening/reshaping them.
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float): learning rate (default: 1e-3).
+        beta1 (float): momentum factor (default: 0.9).
+        weight_decay (float): weight decay (L2 penalty) (default: 0).
+        nesterov (bool): enables Nesterov momentum (default: True).
+        ns_steps (int): number of Newton-Schulz iterations to perform (default: 5).
+        ns_eps (float): epsilon for Newton-Schulz normalization stability (default: 1e-7).
+        ns_coeffs (tuple[float, float, float]): The (a, b, c) coefficients for the
+            quintic polynomial in the Newton-Schulz iteration.
+            (default: (3.4445, -4.7750, 2.0315)).
+        stochastic_rounding (bool): whether to use stochastic rounding for
+            BF16 parameter updates (default: True).
+        vector_reshape_muon (bool): whether to reshape 1D vectors into 2D
+            matrices for muon NewtonSchulz (default: False).
+        vector_reshape (bool): whether to reshape 1D vectors into 2D
+            matrices to apply low-rank compression (default: True).
+        nnmf_factor (bool): whether to use the factorization or disable it to use
+            the uncompressed optimizer. (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-3,
+        beta1: float = 0.9,
+        weight_decay: float = 0.0,
+        nesterov: bool = True,
+        ns_steps: int = 5,
+        ns_eps: float = 1e-7,
+        ns_coeffs: tuple[float, float, float] = (3.4445, -4.7750, 2.0315),
+        stochastic_rounding: bool = True,
+        vector_reshape_muon: bool = False,
+        vector_reshape: bool = True,
+        nnmf_factor: bool = False,
+    ):
+        if not (lr >= 0.0):
+            raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
+        if not (0.0 <= beta1 < 1.0):
+            raise ValueError(f"beta1 should be in [0.0, 1.0). Got {beta1}")
+        if not (weight_decay >= 0.0):
+            raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
+        if not (ns_steps > 0):
+            raise ValueError(f"Newton-Schulz steps should be > 0. Got {ns_steps}")
+
+        defaults = {
+            "lr": lr, "beta1": beta1, "weight_decay": weight_decay,
+            "nesterov": nesterov, "ns_steps": ns_steps, "ns_eps": ns_eps,
+            "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor,
+            "vector_reshape": vector_reshape,
+            "vector_reshape_muon": vector_reshape_muon,
+        }
+        self.stochastic_rounding = stochastic_rounding
+        super().__init__(params, defaults)
+
+    @property
+    def supports_fused_back_pass(self):
+        return True
+
+    @property
+    def supports_memory_efficient_fp16(self):
+        return True
+
+    @property
+    def supports_flat_params(self):
+        return False
+
+    @torch.no_grad()
+    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
+        if p.grad is None:
+            return
+
+        grad = p.grad
+        state = self.state[p]
+
+        # State Initialization
+        if 'step' not in state:
+            state['step'] = 0
+
+            should_factor = (
+                group['nnmf_factor'] and
+                not (len(p.shape) == 1 and not group['vector_reshape'])
+            )
+
+            state['factored'] = should_factor
+
+            state['reshaped_1d_muon'] = len(p.shape) == 1 and group['vector_reshape_muon']
+
+            dtype = torch.float32 if group['nnmf_factor'] else p.dtype
+            device = p.device
+            if group['vector_reshape'] or state['reshaped_1d_muon']:
+                    state['effective_shape'] = _get_effective_shape(p.numel())
+                    d1, d2 = state['effective_shape']
+            if state['factored']:
+                    state['mu_m_nmf'] = torch.zeros(d1, device=device, dtype=dtype) 
+                    state['mv_m_nmf'] = torch.zeros(d2, device=device, dtype=dtype)
+                    packed_d2 = (d2 + 7) // 8
+                    state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=device)
+            else:
+                if len(p.shape) >= 2:
+                    state['momentum_buffer'] = torch.zeros_like(p)
+                if state['reshaped_1d_muon']:
+                    state['momentum_buffer'] = torch.zeros((d1, d2), device=device, dtype=dtype)
+                elif len(p.shape) == 1:
+                    state['momentum_buffer'] = torch.zeros_like(p)
+
+        beta1 = group['beta1']
+        nesterov = group['nesterov']
+
+        if state['factored']: # Factored Muon
+
+            # Reconstruct momentum from previous step's factors & sign
+            d1, d2 = state['effective_shape']
+            mt_buf = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
+            unpacked_sign = _unpack_bools(state['sign'], original_m=d2)
+            torch.where(unpacked_sign, mt_buf, -mt_buf, out=mt_buf)
+            del unpacked_sign
+
+            # Update momentum in full-size
+            grad_reshaped = grad.view(d1, d2)
+            mt_buf.mul_(beta1).add_(grad_reshaped)
+
+            if nesterov:
+                # Nesterov momentum
+                update = grad_reshaped.add(mt_buf, alpha=beta1)
+            else:
+                # Standard momentum
+                update = mt_buf.clone()
+            del grad_reshaped
+
+            update = _newton_schulz_iteration(
+                update,
+                steps=group['ns_steps'],
+                eps=group['ns_eps'],
+                coeffs=group['ns_coeffs'],
+            )
+
+            update = update.view(p.shape).mul_(group['lr'])
+
+            state['sign'] = _pack_bools(mt_buf > 0)
+            _nnmf(mt_buf.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
+            del mt_buf
+
+        else: # Standard Muon logic for non-factored tensors
+
+            if len(p.shape) >= 2 or state['reshaped_1d_muon']:
+
+                # Momentum update
+                mt_buf = state['momentum_buffer']
+                if state['reshaped_1d_muon']:
+                    d1, d2 = state['effective_shape']
+                    grad_reshaped = grad.view(d1, d2)
+                    mt_buf.mul_(beta1).add_(grad_reshaped)
+                else:
+                    mt_buf.mul_(beta1).add_(grad)
+
+                if nesterov:
+                    # Nesterov momentum
+                    if state['reshaped_1d_muon']:
+                        update = grad_reshaped.add(mt_buf, alpha=beta1)
+                        del grad_reshaped
+                    else:
+                        update = grad.add(mt_buf, alpha=beta1)
+                else:
+                    # Standard momentum
+                    update = mt_buf.clone()
+
+                # For Conv layers (4D) or other high-dim tensors, flatten to 2D
+                if len(p.shape) > 2:
+                    update = update.view(p.shape[0], -1)
+
+                # NewtonSchulz
+                update = _newton_schulz_iteration(
+                    update,
+                    steps=group['ns_steps'],
+                    eps=group['ns_eps'],
+                    coeffs=group['ns_coeffs'],
+                )
+
+                # Reshape back to original if we flattened or reshaped
+                if len(p.shape) > 2 or state['reshaped_1d_muon']:
+                    update = update.view(p.shape)
+
+                update.mul_(group['lr'])
+            
+            else: # Fallback to standard SGD with momentum for 1D params (biases, etc.) when not reshaped
+                # Momentum update
+                mt_buf = state['momentum_buffer']
+                mt_buf.mul_(beta1).add_(grad)
+                if nesterov:
+                    # Nesterov momentum
+                    update = grad.add(mt_buf, alpha=beta1)
+                else:
+                    # Standard momentum
+                    update = mt_buf.clone()
+                update.mul_(group['lr'])
+
+        # Decoupled weight decay
+        if group["weight_decay"] != 0:
+            if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+                add_stochastic_(p.data, p.data, alpha=-group["weight_decay"] * group["lr"])
+            else:
+                p.data.add_(p.data, alpha=-group["weight_decay"] * group["lr"])
+
+        if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+            add_stochastic_(p.data, -update)
+        else:
+            p.data.add_(-update)
+        del update
+
+        state['step'] += 1
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for i, p in enumerate(group['params']):
+                self.step_parameter(p, group, i)
+
+        return loss

adv_optm/optim/Prodigy_adv.py

@@ -304,7 +304,7 @@ class Prodigy_adv(torch.optim.Optimizer):
                 state['p0'] = torch.tensor(0, device=device, dtype=p.dtype)
 
         current_step = state['step']
-        if group['kourkoutas_beta']:
+        if group.get('kourkoutas_beta', False):
             # Call prepare_step() once at the beginning of the step for all params
             self.kourkoutas_helper.maybe_prepare_step(current_step)
             # Accumulate current grad's norm for the *next* step
@@ -515,7 +515,7 @@ class Prodigy_adv(torch.optim.Optimizer):
             d_hat = self.d
             if global_d_denom > 0:
                 d_hat = d_coef * global_d_numerator / global_d_denom
-                if g_group['d_limiter']:
+                if g_group.get('d_limiter', False):
                     d_hat = min(self.d * (2 ** 0.25), d_hat)
                 if self.d == g_group['d0']:
                     self.d = max(self.d, d_hat)

adv_optm/optim/Simplified_AdEMAMix.py

@@ -191,7 +191,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
         beta1_final, beta2 = group["betas"]
 
         current_step = state['step']
-        if group['kourkoutas_beta']:
+        if group.get('kourkoutas_beta', False):
             # Call prepare_step() once at the beginning of the step for all params
             self.kourkoutas_helper.maybe_prepare_step(current_step)
             # Accumulate current grad's norm for the *next* step
@@ -210,7 +210,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
 
         if group['use_bias_correction']:
             state['num_sum'] = beta1 * state['num_sum'] + 1.0
-            if group['kourkoutas_beta']:
+            if group.get('kourkoutas_beta', False):
                 state['den_sum'] = group['betas'][1] * state['den_sum'] + (1.0 - group['betas'][1])
             else:
                 state['den_sum'] = beta2 * state['den_sum'] + (1.0 - beta2)

adv_optm/optim/__init__.py

@@ -4,6 +4,7 @@ from .Adopt_adv import Adopt_adv
 from .Simplified_AdEMAMix import Simplified_AdEMAMix
 from .Lion_adv import Lion_adv
 from .Lion_Prodigy_adv import Lion_Prodigy_adv
+from .Muon_adv import Muon_adv
 
 __all__ = [
     "AdamW_adv",
@@ -12,4 +13,5 @@ __all__ = [
     "Simplified_AdEMAMix",
     "Lion_adv",
     "Lion_Prodigy_adv",
+    "Muon_adv",
 ]

adv_optm/util/Newton_Schulz.py

@@ -0,0 +1,48 @@
+import torch
+
+@torch.no_grad()
+def _newton_schulz_iteration(
+    G: torch.Tensor,
+    steps: int = 5,
+    eps: float = 1e-7,
+    coeffs: tuple[float, float, float] = (3.4445, -4.7750, 2.0315)
+) -> torch.Tensor:
+    """
+    Performs the Newton-Schulz iteration to find the nearest orthogonal matrix.
+    This is the core computation of the Muon optimizer.
+
+    Args:
+        G (torch.Tensor): The 2D input matrix (momentum-accumulated gradient).
+        steps (int): The number of iterations to run.
+        eps (float): Small constant for numerical stability during normalization.
+        coeffs (tuple[float, float, float]): The (a, b, c) coefficients for the
+            quintic polynomial update.
+
+    Returns:
+        torch.Tensor: The orthogonalized matrix.
+    """
+    assert G.ndim == 2, "Newton-Schulz iteration only supports 2D matrices."
+
+    a, b, c = coeffs
+
+    X = G.to(torch.bfloat16)
+
+    # Normalize the matrix
+    X.div_(X.norm() + eps)
+
+    # Handle non-square matrices by transposing the taller one
+    transposed = G.size(0) > G.size(1)
+    if transposed:
+        X = X.T
+
+    # Perform the iterative updates
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * (A @ A)
+        X = a * X + B @ X
+
+    # Transpose back if necessary
+    if transposed:
+        X = X.T
+
+    return X.to(G.dtype)

adv_optm/util/__init__.py

@@ -2,10 +2,11 @@ from .BF16_Stochastic_Rounding import add_stochastic_
 from .Effective_Shape import _get_effective_shape
 from .One_Bit_Boolean import _pack_bools, _unpack_bools
 from .OrthoGrad import _orthogonalize_gradient
-
+from .Newton_Schulz import _newton_schulz_iteration
 __all__ = [
     "_pack_bools", "_unpack_bools",
     "add_stochastic_",
     "_get_effective_shape",
     "_orthogonalize_gradient",
+    "_newton_schulz_iteration",
 ]

{adv_optm-1.1.2.dist-info → adv_optm-1.2.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 1.1.2
+Version: 1.2.dev1
 Summary: A family of highly efficient, lightweight yet powerful optimizers.
 Home-page: https://github.com/Koratahiu/Advanced_Optimizers
 Author: Koratahiu

adv_optm-1.2.dev1.dist-info/RECORD

@@ -0,0 +1,22 @@
+adv_optm/__init__.py,sha256=fdpoxqapAZbMiPax_P4Zm9PkN--71G0yds0q0V9oAbo,341
+adv_optm/optim/AdamW_adv.py,sha256=7vWfPS2J54U9ZKFQiNJ_l86PvITb0MQ61Fy4Fzmf1d4,17479
+adv_optm/optim/Adopt_adv.py,sha256=NXbtPrGm3tZr06cApi5oEHZ2F1zwss3tRi15SGnrYPc,21426
+adv_optm/optim/Lion_Prodigy_adv.py,sha256=LEA3UYJpPeFnmxeniLNv1u2LKKj4ufx3Bq_MLw-nWXk,14617
+adv_optm/optim/Lion_adv.py,sha256=aGNAplZlyXYgVllYcV_s4bK8iC4fv6EizFoWIMNLdBc,8299
+adv_optm/optim/Muon_adv.py,sha256=eXqPL6GIWutBJpP7Yb_qIk7pGAjwfTAloCFRDhkRoUU,9908
+adv_optm/optim/Prodigy_adv.py,sha256=0_XG5YnMQTv-zJysJHlJniSo5kGYdX3p3o1e33HLt78,25897
+adv_optm/optim/Simplified_AdEMAMix.py,sha256=nEIA3yM11nBooKzHudB5l3x4UdFRBYRwiKVUkGmO0K8,12971
+adv_optm/optim/__init__.py,sha256=3o2XJ4J-PUq3rJM2mBnmuHwbKNb4LuW-Ig_9aBC0ycc,431
+adv_optm/util/BF16_Stochastic_Rounding.py,sha256=Q5H0BcogmE4atP65dLoI21HKSf50lRdsBDfeF6v9Tbg,1548
+adv_optm/util/Effective_Shape.py,sha256=TBvIk1V8IuTbbBsxuekJA4e_v8JlR5Nujtut8RTWAm4,318
+adv_optm/util/Kourkoutas.py,sha256=woyJfX7l4eieeg0pC5XrILBLvwECwbD3a6ou1K6qjKU,8706
+adv_optm/util/NNMF.py,sha256=yRf5IP5Sjq0Uf0DxN0Q8NxEGSdD-f1ULziLVDOjY8K4,639
+adv_optm/util/Newton_Schulz.py,sha256=wJ_sKRaGVIsOofQ737my4ng494qX_pfgOqlDDmYtnCg,1377
+adv_optm/util/One_Bit_Boolean.py,sha256=Wat49esdwohuN-OHOFMW8D0aOQgV9cP5Rl8z6yfmpos,1068
+adv_optm/util/OrthoGrad.py,sha256=NzInuBQGy_Ja__M1R9XbvqVaQ0fhGbtGgFE9YON7B3I,707
+adv_optm/util/__init__.py,sha256=jAaUfaAjFrTJ6-Q915ezAbq0efRbpYjriW2OdeCbSzo,433
+adv_optm-1.2.dev1.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-1.2.dev1.dist-info/METADATA,sha256=ofbAQu0ldYk8udMEC0jLcI9Ex2a6M8iaXheDTo4Un3M,14022
+adv_optm-1.2.dev1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+adv_optm-1.2.dev1.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-1.2.dev1.dist-info/RECORD,,

adv_optm-1.1.2.dist-info/RECORD

@@ -1,20 +0,0 @@
-adv_optm/__init__.py,sha256=IJAqLP1mOIBeEFFeMCrFvxEwK7oz-g5SRAEfmukmy9o,306
-adv_optm/optim/AdamW_adv.py,sha256=ddEUVOif1gfZPgEJNrEGZ2wnha4MPMWw5ppPd8acQ3o,17457
-adv_optm/optim/Adopt_adv.py,sha256=fhH3hS9K6z5Blxc7NFfzpCrUGbl9EQnwLPmKDxBC1zg,21415
-adv_optm/optim/Lion_Prodigy_adv.py,sha256=0AY-2hjdNKnxiY4MUG4Y4HCe1DvesuwaaxRzrHUkAGA,14606
-adv_optm/optim/Lion_adv.py,sha256=aGNAplZlyXYgVllYcV_s4bK8iC4fv6EizFoWIMNLdBc,8299
-adv_optm/optim/Prodigy_adv.py,sha256=nD59cAWOJJCjZdIiuD5hD9MWO5sTjPQSvq-3dwGTcEM,25875
-adv_optm/optim/Simplified_AdEMAMix.py,sha256=gPjMhKulzmAeO42foe-d7xW0AcB50vKFYsvHgxbD3uc,12949
-adv_optm/optim/__init__.py,sha256=pcP865H2j1tut2VfTUhzQh7V8TF_tzPjqFnjMfFed2k,382
-adv_optm/util/BF16_Stochastic_Rounding.py,sha256=Q5H0BcogmE4atP65dLoI21HKSf50lRdsBDfeF6v9Tbg,1548
-adv_optm/util/Effective_Shape.py,sha256=TBvIk1V8IuTbbBsxuekJA4e_v8JlR5Nujtut8RTWAm4,318
-adv_optm/util/Kourkoutas.py,sha256=woyJfX7l4eieeg0pC5XrILBLvwECwbD3a6ou1K6qjKU,8706
-adv_optm/util/NNMF.py,sha256=yRf5IP5Sjq0Uf0DxN0Q8NxEGSdD-f1ULziLVDOjY8K4,639
-adv_optm/util/One_Bit_Boolean.py,sha256=Wat49esdwohuN-OHOFMW8D0aOQgV9cP5Rl8z6yfmpos,1068
-adv_optm/util/OrthoGrad.py,sha256=NzInuBQGy_Ja__M1R9XbvqVaQ0fhGbtGgFE9YON7B3I,707
-adv_optm/util/__init__.py,sha256=qoyIF0jcLjs_vSEcsv36clw5LFNBEbifyXrrVxMH-G4,349
-adv_optm-1.1.2.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-1.1.2.dist-info/METADATA,sha256=mtTfygEQn52Jqwc_W7rnDhJwRAArvaEjiK0s-cyDFVQ,14019
-adv_optm-1.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-1.1.2.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-1.1.2.dist-info/RECORD,,