adv-optm 2.4.dev11__tar.gz → 2.4.dev13__tar.gz

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/PKG-INFO

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

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/adv_optm/__init__.py

@@ -8,7 +8,7 @@ from .optim import (
     Muon_adv,
     AdaMuon_adv,
     SignSGD_adv,
-    SGD_adv,
+    SinkSGD_adv,
 )
 
 __all__ = [
@@ -21,7 +21,7 @@ __all__ = [
     "Muon_adv",
     "AdaMuon_adv",
     "SignSGD_adv",
-    "SGD_adv",
+    "SinkSGD_adv",
 ]
 
-__version__ = "2.4.dev11"
+__version__ = "2.4.dev13"

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/adv_optm/optim/AdaMuon_adv.py

@@ -468,7 +468,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 if actual_precision == 'bf16_sr' and random_int_state_tensor is not None:
                     random_int_state_tensor = param_update._get_random_int_for_sr(p)
                 elif actual_precision == 'int8_sr':
-                    random_int_state_tensor = param_update._get_random_int_for_int8_sr(p)
+                    random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
                 elif actual_precision == 'fp8_sr':
                     random_int_state_tensor = param_update._get_random_int_for_fp8_sr(p)
             else:

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/adv_optm/optim/AdamW_adv.py

@@ -232,12 +232,13 @@ class AdamW_adv(torch.optim.Optimizer):
     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
+    def init_step(self):
+        for group in self.param_groups:
+            for i, p in enumerate(group['params']):
+                self.__init_state(p, group)
 
-        grad = p.grad
+    @torch.no_grad()
+    def __init_state(self, p, group):
         state = self.state[p]
 
         # State Initialization
@@ -303,6 +304,15 @@ class AdamW_adv(torch.optim.Optimizer):
 
             _init_fisher_wd_scaler(group, state, p)
 
+    @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]
+        self.__init_state(p, group)
+
         beta1, beta2 = group['betas']
 
         current_step = state['step']
@@ -333,7 +343,7 @@ class AdamW_adv(torch.optim.Optimizer):
             if group['actual_state_precision'] == 'bf16_sr' and random_int_state_tensor is None:
                 random_int_state_tensor = param_update._get_random_int_for_sr(p)
             elif group['actual_state_precision'] == 'int8_sr':
-                random_int_state_tensor = param_update._get_random_int_for_int8_sr(p)
+                random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
             elif group['actual_state_precision'] == 'fp8_sr':
                 random_int_state_tensor = param_update._get_random_int_for_fp8_sr(p)
             step_param_fn = self._compiled_step_parameter

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/adv_optm/optim/Adopt_adv.py

@@ -359,7 +359,7 @@ class Adopt_adv(torch.optim.Optimizer):
             if group['actual_state_precision'] == 'bf16_sr' and random_int_state_tensor is None:
                 random_int_state_tensor = param_update._get_random_int_for_sr(p)
             elif group['actual_state_precision'] == 'int8_sr':
-                random_int_state_tensor = param_update._get_random_int_for_int8_sr(p)
+                random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
             elif group['actual_state_precision'] == 'fp8_sr':
                 random_int_state_tensor = param_update._get_random_int_for_fp8_sr(p)
             step_param_fn = self._compiled_step_parameter

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/adv_optm/optim/Muon_adv.py

@@ -422,7 +422,7 @@ class Muon_adv(torch.optim.Optimizer):
                 if actual_precision == 'bf16_sr' and random_int_state_tensor is not None:
                     random_int_state_tensor = param_update._get_random_int_for_sr(p)
                 elif actual_precision == 'int8_sr':
-                    random_int_state_tensor = param_update._get_random_int_for_int8_sr(p)
+                    random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
                 elif actual_precision == 'fp8_sr':
                     random_int_state_tensor = param_update._get_random_int_for_fp8_sr(p)
             else:

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/adv_optm/optim/SignSGD_adv.py

@@ -252,7 +252,7 @@ class SignSGD_adv(torch.optim.Optimizer):
                 if group['actual_state_precision'] == 'bf16_sr' and random_int_state_tensor is None:
                     random_int_state_tensor = param_update._get_random_int_for_sr(p)
                 elif group['actual_state_precision'] == 'int8_sr':
-                    random_int_state_tensor = param_update._get_random_int_for_int8_sr(p)
+                    random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
                 elif group['actual_state_precision'] == 'fp8_sr':
                     random_int_state_tensor = param_update._get_random_int_for_fp8_sr(p)
 

adv_optm-2.4.dev11/adv_optm/optim/SGD_adv.py → adv_optm-2.4.dev13/adv_optm/optim/SinkSGD_adv.py

@@ -11,12 +11,11 @@ from ..util.centered_decay import _init_anchor
 from ..util.state_util import init_state_tensor, get_state, set_state, upcast_grad_for_precision
 from ..util.sinkhorn import apply_sr_sinkhorn
 
-class SGD_adv(torch.optim.Optimizer):
+class SinkSGD_adv(torch.optim.Optimizer):
     """
-    Implements an advanced Stochastic Gradient Descent (SGD) algorithm.
-    This is an advanced version of SGD with optional features like
-    low-rank factorization of optimizer states (SMMF), OrthoGrad, 
-    Cautious updating, and AdEMAMix extensions.
+    Implements an advanced Stochastic Gradient Descent (SGD) with Sinkhorn Iterative Normalization (SinkSGD) algorithm.
+    This is an advanced version of SinkSGD with optional features like
+    low-rank factorization of optimizer states (SMMF), OrthoGrad, etc.
 
     Args:
         params (iterable): iterable of parameters to optimize or dicts defining
@@ -62,11 +61,11 @@ class SGD_adv(torch.optim.Optimizer):
         cautious_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
-        # OrthoGrad
-        orthogonal_gradient: bool = False,
         # Sinkhorn Iterative Normalization
-        sinkhorn: bool = False,
         sinkhorn_iterations: int = 5,
+        orthogonal_sinkhorn: bool = False,
+        # OrthoGrad
+        orthogonal_gradient: bool = False,
         # Spectral Normed Optimizer
         spectral_normalization: bool = False,
         # Centered WD
@@ -101,7 +100,8 @@ class SGD_adv(torch.optim.Optimizer):
             "decoupled_wd": decoupled_wd, "cautious_wd": cautious_wd,
             "orthogonal_gradient": orthogonal_gradient, 
             "compiled_optimizer": compiled_optimizer,
-            "sinkhorn": sinkhorn, "sinkhorn_iterations": sinkhorn_iterations,
+            "sinkhorn_iterations": sinkhorn_iterations,
+            "orthogonal_sinkhorn": orthogonal_sinkhorn,
             "spectral_normalization": spectral_normalization,
             "centered_wd": centered_wd, "centered_wd_mode": centered_wd_mode,
             "state_precision": state_precision,
@@ -116,6 +116,8 @@ class SGD_adv(torch.optim.Optimizer):
             for device in devices:
                 param_update.set_seed(device)
 
+        self.init_step()
+
         self._compiled_step_parameter = None
         if compiled_optimizer:
             self.compile(fullgraph=True)
@@ -136,14 +138,14 @@ class SGD_adv(torch.optim.Optimizer):
     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
+    def init_step(self):
+        for group in self.param_groups:
+            for i, p in enumerate(group['params']):
+                self.__init_state(p, group)
 
-        grad = p.grad
+    @torch.no_grad()
+    def __init_state(self, p, group):
         state = self.state[p]
-
         # State Initialization
         if 'step' not in state:
             state['step'] = 0
@@ -180,6 +182,15 @@ class SGD_adv(torch.optim.Optimizer):
 
             _init_anchor(p, state, group)
 
+    @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]
+        self.__init_state(p, group)
+
         step_size = group['lr']
 
         random_int_tensor = None
@@ -193,7 +204,7 @@ class SGD_adv(torch.optim.Optimizer):
             if group['actual_state_precision'] == 'bf16_sr' and random_int_state_tensor is None:
                 random_int_state_tensor = param_update._get_random_int_for_sr(p)
             elif group['actual_state_precision'] == 'int8_sr':
-                random_int_state_tensor = param_update._get_random_int_for_int8_sr(p)
+                random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
             elif group['actual_state_precision'] == 'fp8_sr':
                 random_int_state_tensor = param_update._get_random_int_for_fp8_sr(p)
             step_param_fn = self._compiled_step_parameter
@@ -219,7 +230,7 @@ class SGD_adv(torch.optim.Optimizer):
 
             if momentum != 0:
                 buf = _reconstruct_state((state['mu_b_nmf'], state['mv_b_nmf'], state['sign'], d2), signed=True)
-                buf.mul_(momentum).add_(grad_reshaped, alpha=1 - momentum)
+                buf.lerp_(grad_reshaped, 1 - momentum)
 
                 # Factorize updated buffer
                 state['mu_b_nmf'], state['mv_b_nmf'], state['sign'] = _factorize_state(buf.clone(), signed=True)
@@ -239,9 +250,7 @@ class SGD_adv(torch.optim.Optimizer):
 
             if momentum != 0:
                 buf = get_state(state, 'momentum_buffer', actual_precision)
-
-                buf.mul_(momentum).add_(grad, alpha=1 - momentum)
-
+                buf.lerp_(grad, 1 - momentum)
 
                 set_state(state, 'momentum_buffer', buf, actual_precision, random_int_state_tensor)
 
@@ -254,8 +263,8 @@ class SGD_adv(torch.optim.Optimizer):
 
             del random_int_state_tensor
 
-        if group['sinkhorn']:
-            update = apply_sr_sinkhorn(update, iters=group['sinkhorn_iterations'])
+        # Sinkhorn iterative normalization
+        update = apply_sr_sinkhorn(update, p, ortho_project=group['orthogonal_sinkhorn'], iters=group['sinkhorn_iterations'])
 
         update_scaling = step_size
         if group.get('spectral_normalization', False):

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/adv_optm/optim/__init__.py

@@ -7,7 +7,7 @@ from .Lion_Prodigy_adv import Lion_Prodigy_adv
 from .Muon_adv import Muon_adv
 from .AdaMuon_adv import AdaMuon_adv
 from .SignSGD_adv import SignSGD_adv
-from .SGD_adv import SGD_adv
+from .SinkSGD_adv import SinkSGD_adv
 
 __all__ = [
     "AdamW_adv",
@@ -19,5 +19,5 @@ __all__ = [
     "Muon_adv",
     "AdaMuon_adv",
     "SignSGD_adv",
-    "SGD_adv",
+    "SinkSGD_adv",
 ]

adv_optm-2.4.dev13/adv_optm/util/sinkhorn.py

@@ -0,0 +1,77 @@
+import math
+import torch
+
+def apply_sr_sinkhorn(update: torch.Tensor, p: torch.Tensor, ortho_project: bool, iters: int = 5) -> torch.Tensor:
+    """
+    Applies Square-Root Sinkhorn (SR-Sinkhorn) multi-normalization.
+    As described in 'Gradient Multi-Normalization for Efficient LLM Training'.
+
+    This technique normalizes a 2D matrix alternatively by its row-wise L2 norm 
+    and column-wise L2 norm, driving it toward a fixed point that uniformly 
+    distributes update magnitudes.
+    """
+    original_shape = update.shape
+    original_dtype = update.dtype
+    update = update.float()
+
+    # 1D Vector Case
+    if update.dim() == 1:
+        if ortho_project:
+            p_float = p.float()
+            p_norm_sq = torch.dot(p_float, p_float).add_(1e-30)
+            proj = torch.dot(p_float, update) / p_norm_sq
+            update.sub_(p_float * proj) 
+        norm = update.norm(p=2).clamp_min_(1e-12)
+        return update.mul_(math.sqrt(update.numel()) / norm).view(original_shape).to(original_dtype)
+
+    # 2D+ Matrix Case
+    update_2d = update.view(update.shape[0], -1)
+
+    m, n = update_2d.shape
+
+    # Dynamically determine the order of normalization based on aspect ratio
+    # Normalizing the longer dimension first aids stability.
+    scale_cond = update_2d.shape[0] > update_2d.shape[1]
+    dim = 0 if scale_cond else 1
+
+
+    # Precompute scaling factors. 
+    scale_first = math.sqrt(m if scale_cond else n)
+    scale_second = math.sqrt(n if scale_cond else m)
+
+    if ortho_project:
+        param_2d = p.float().view(p.shape[0], -1)
+        p_norm_sq_dim = torch.sum(param_2d * param_2d, dim=dim, keepdim=True).add_(1e-30)
+        p_norm_sq_adim = torch.sum(param_2d * param_2d, dim=1-dim, keepdim=True).add_(1e-30)
+
+    # In-place alternating Sinkhorn normalization steps
+    for _ in range(iters):
+        # First normalization step
+        norm1 = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(1e-12)
+        update_2d.mul_(scale_first / norm1)
+        if ortho_project:
+            update_2d = ortho_normed(param_2d, update_2d, p_norm_sq_dim, dim, scale_first)
+
+        # Second normalization step
+        norm2 = update_2d.norm(p=2, dim=1-dim, keepdim=True).clamp_min_(1e-12)
+        update_2d.mul_(scale_second / norm2)
+        if ortho_project:
+            update_2d = ortho_normed(param_2d, update_2d, p_norm_sq_adim, 1-dim, scale_second)
+
+    return update_2d.view(original_shape).to(original_dtype)
+
+def ortho_normed(p_2d, update_2d, p_norm_sq, dim, target_norm):
+    """
+    Projects the update to be orthogonal to p along 'dim' and restores the original norm.
+    """
+    # Project: g_orth = g - (p * <p, g> / ||p||^2)
+    dot_prod = torch.sum(p_2d * update_2d, dim=dim, keepdim=True)
+    proj = dot_prod / p_norm_sq
+
+    # In-place subtraction: update_2d = update_2d - (proj * p_2d)
+    update_2d.addcmul_(proj, p_2d, value=-1.0)
+
+    # Magnitude Preservation
+    g_orth_norm = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(1e-12)
+    scale_factor = target_norm / g_orth_norm
+    return update_2d.mul_(scale_factor)

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/adv_optm.egg-info/SOURCES.txt

@@ -14,9 +14,9 @@ adv_optm/optim/Lion_Prodigy_adv.py
 adv_optm/optim/Lion_adv.py
 adv_optm/optim/Muon_adv.py
 adv_optm/optim/Prodigy_adv.py
-adv_optm/optim/SGD_adv.py
 adv_optm/optim/SignSGD_adv.py
 adv_optm/optim/Simplified_AdEMAMix.py
+adv_optm/optim/SinkSGD_adv.py
 adv_optm/optim/__init__.py
 adv_optm/util/Kourkoutas.py
 adv_optm/util/Muon_AuxAdam.py

{adv_optm-2.4.dev11 → adv_optm-2.4.dev13}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="adv_optm",
-    version="2.4.dev11",
+    version="2.4.dev13",
     author="Koratahiu",
     author_email="hiuhonor@gmail.com",
     license='Apache 2.0',

adv_optm-2.4.dev11/adv_optm/util/sinkhorn.py

@@ -1,42 +0,0 @@
-import math
-import torch
-
-def apply_sr_sinkhorn(update: torch.Tensor, iters: int = 5) -> torch.Tensor:
-    """
-    Applies Square-Root Sinkhorn (SR-Sinkhorn) multi-normalization.
-    As described in 'Gradient Multi-Normalization for Efficient LLM Training'.
-
-    This technique normalizes a 2D matrix alternatively by its row-wise L2 norm 
-    and column-wise L2 norm, driving it toward a fixed point that uniformly 
-    distributes update magnitudes.
-    """
-    original_shape = update.shape
-
-    if update.dim() == 1:
-            norm = update.norm(p=2).clamp_min_(1e-12)
-            return update.mul_(math.sqrt(update.numel()) / norm)
-    else:
-        # Flatten >= 3D tensors into 2D matrices
-        update_2d = update.view(update.shape[0], -1)
-
-    m, n = update_2d.shape
-
-    # Dynamically determine the order of normalization based on aspect ratio
-    # Normalizing the longer dimension first aids stability.
-    dim = 0 if m > n else 1
-
-    # Precompute scaling factors. 
-    scale_first = math.sqrt(m) if dim == 0 else math.sqrt(n)
-    scale_second = math.sqrt(n) if dim == 0 else math.sqrt(m)
-
-    # In-place alternating Sinkhorn normalization steps
-    for _ in range(iters):
-        # First normalization step
-        norm1 = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(1e-12)
-        update_2d.mul_(scale_first / norm1)
-
-        # Second normalization step
-        norm2 = update_2d.norm(p=2, dim=1-dim, keepdim=True).clamp_min_(1e-12)
-        update_2d.mul_(scale_second / norm2)
-
-    return update_2d.view(original_shape)