adv-optm 2.6.1.dev3__tar.gz → 2.6.2.dev2__tar.gz

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/PKG-INFO

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

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "SinkSGD_adv",
 ]
 
-__version__ = "2.6.1.dev3"
+__version__ = "2.6.2.dev2"

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/optim/AdaMuon_adv.py

@@ -137,7 +137,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
         # Decoupled/cautious weight decay
         weight_decay: float = 0,
         cautious_wd: bool = False,
-        scaled_wd: bool = False,
         # Nesterov momentum
         nesterov: bool = True,
         nesterov_coef: float | None = None,
@@ -228,7 +227,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
         defaults = {
             "lr": lr, "betas": betas, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
-            "eps": eps, "rms_rescaling": rms_rescaling, "ns_steps": ns_steps, "scaled_wd": scaled_wd,
+            "eps": eps, "rms_rescaling": rms_rescaling, "ns_steps": ns_steps,
             "ns_eps": ns_eps, "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor,
             "vector_reshape": vector_reshape,
             "nesterov":nesterov, "nesterov_coef": nesterov_coef, "use_atan2":use_atan2,

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/optim/AdamW_adv.py

@@ -98,7 +98,6 @@ class AdamW_adv(torch.optim.Optimizer):
         weight_decay: float = 0.0,
         fisher_wd: bool = False,
         cautious_wd: bool = False,
-        scaled_wd: bool = False,
         # Adam's Bias Correction
         use_bias_correction: bool = True,
         # Stochastic Rounding for BF16
@@ -157,7 +156,7 @@ class AdamW_adv(torch.optim.Optimizer):
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
-            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, "scaled_wd": scaled_wd,
+            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
             "use_atan2": use_atan2, "nesterov": nesterov, "nesterov_coef": nesterov_coef,
             "normed_momentum": normed_momentum,
             "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction,

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/optim/Adopt_adv.py

@@ -101,7 +101,6 @@ class Adopt_adv(torch.optim.Optimizer):
         weight_decay: float = 0.0,
         fisher_wd: bool = False,
         cautious_wd: bool = False,
-        scaled_wd: bool = False,
         # ADOPT clipping
         clip_lambda: Optional[Callable[[int], float]] = lambda step: step**0.25,
         # Adam_atan2 (scale invariant)
@@ -158,7 +157,7 @@ class Adopt_adv(torch.optim.Optimizer):
             state_precision = "factored"
 
         defaults = {
-            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "scaled_wd": scaled_wd,
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
             "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, "orthogonal_gradient": orthogonal_gradient,
             "nesterov": nesterov, "nesterov_coef": nesterov_coef,
             "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
@@ -315,7 +314,8 @@ class Adopt_adv(torch.optim.Optimizer):
         # The first step is for initialization only (skip when use_atan2 as it's scale invariant).
         if state['step'] == 0 and not (self.use_atan2 or group.get('spectral_normalization', False)):
             state['step'] += 1
-            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+            if group.get('kourkoutas_beta', False):
+                self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
             return
 
         random_int_tensor = None

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/optim/Lion_adv.py

@@ -64,7 +64,6 @@ class Lion_adv(torch.optim.Optimizer):
         # Decoupled/cautious weight decay
         weight_decay: float = 0.0,
         cautious_wd: bool = False,
-        scaled_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
@@ -97,7 +96,6 @@ class Lion_adv(torch.optim.Optimizer):
             betas=betas,
             weight_decay=weight_decay,
             cautious_wd=cautious_wd,
-            scaled_wd=scaled_wd,
             vector_reshape=vector_reshape,
             orthogonal_gradient=orthogonal_gradient,
             kappa_p=kappa_p,

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/optim/Muon_adv.py

@@ -111,7 +111,6 @@ class Muon_adv(torch.optim.Optimizer):
         # Decoupled/cautious weight decay
         weight_decay: float = 0.0,
         cautious_wd: bool = False,
-        scaled_wd: bool = False,
         # Nesterov momentum
         nesterov: bool = True,
         nesterov_coef: float | None = None,
@@ -202,7 +201,7 @@ class Muon_adv(torch.optim.Optimizer):
         defaults = {
             "lr": lr, "beta1": beta1, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
             "nesterov": nesterov, "nesterov_coef": nesterov_coef, "ns_steps": ns_steps, "ns_eps": ns_eps,
-            "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor, "scaled_wd": scaled_wd,
+            "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor,
             "vector_reshape": vector_reshape,  "rms_rescaling": rms_rescaling,
             "orthogonal_gradient": orthogonal_gradient,
             'compiled_optimizer': compiled_optimizer,

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/optim/Prodigy_adv.py

@@ -115,7 +115,6 @@ class Prodigy_adv(torch.optim.Optimizer):
         weight_decay: float = 0.0,
         fisher_wd: bool = False,
         cautious_wd: bool = False,
-        scaled_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # Adam_atan2 (scale invariant)
@@ -182,7 +181,7 @@ class Prodigy_adv(torch.optim.Optimizer):
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
-            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, "scaled_wd": scaled_wd,
+            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
             "use_atan2": use_atan2,
             "orthogonal_gradient": orthogonal_gradient,
             "compiled_optimizer": compiled_optimizer,

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/optim/SignSGD_adv.py

@@ -59,7 +59,6 @@ class SignSGD_adv(torch.optim.Optimizer):
         # weight decay features
         geometric_wd: bool = False,
         cautious_wd: bool = False,
-        scaled_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
@@ -109,7 +108,6 @@ class SignSGD_adv(torch.optim.Optimizer):
             momentum=momentum,
             weight_decay=weight_decay,
             cautious_wd=cautious_wd,
-            scaled_wd=scaled_wd,
             geometric_wd=geometric_wd,
             vector_reshape=vector_reshape,
             orthogonal_gradient=orthogonal_gradient,

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/optim/SinkSGD_adv.py

@@ -8,7 +8,7 @@ from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
 from ..util.centered_decay import _init_anchor, dequantize_anchor
 from ..util.state_util import init_state_tensor, get_state, set_state, upcast_grad_for_precision
-from ..util.sinkhorn import apply_sr_sinkhorn, get_sinkhorn_wd_scaler
+from ..util.sinkhorn import apply_sr_sinkhorn, get_sinkhorn_wd_scaler, apply_oft_sinkhorn, apply_spectral_oft_sinkhorn
 from ..util.signed_util import get_signsgd_wd_target
 
 class SinkSGD_adv(torch.optim.Optimizer):
@@ -56,6 +56,8 @@ class SinkSGD_adv(torch.optim.Optimizer):
         # Sinkhorn Iterative Normalization
         sinkhorn_iterations: int = 5,
         orthogonal_sinkhorn: bool = False,
+        # OFT-specific Sinkhorn (symmetric)
+        skew_sinkoft: bool = False,
         # Normalization then Momentum
         normed_momentum: bool = False,
         # SNR Precondition (requires normed_momentum)
@@ -66,7 +68,6 @@ class SinkSGD_adv(torch.optim.Optimizer):
         # weight decay features
         geometric_wd: bool = False,
         cautious_wd: bool = False,
-        scaled_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
@@ -104,11 +105,11 @@ class SinkSGD_adv(torch.optim.Optimizer):
         defaults = {
             "lr": lr, "momentum": momentum,
             "weight_decay": weight_decay, "nesterov": nesterov, "nesterov_coef": nesterov_coef, "normed_momentum": normed_momentum, "snr_cond": snr_cond,
-            "geometric_wd": geometric_wd, "cautious_wd": cautious_wd, "scaled_wd": scaled_wd,
+            "geometric_wd": geometric_wd, "cautious_wd": cautious_wd,
             "orthogonal_gradient": orthogonal_gradient, 
             "compiled_optimizer": compiled_optimizer,
             "sinkhorn_iterations": sinkhorn_iterations,
-            "orthogonal_sinkhorn": orthogonal_sinkhorn,
+            "orthogonal_sinkhorn": orthogonal_sinkhorn, "skew_sinkoft": skew_sinkoft,
             "spectral_normalization": spectral_normalization,
             "centered_wd": centered_wd, "centered_wd_mode": centered_wd_mode,
             "state_precision": state_precision,
@@ -237,6 +238,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
         nesterov = group['nesterov']
         nesterov_coef = group.get('nesterov_coef', None)
         snr_cond = group.get('snr_cond', False)
+        oft_sym = group.get('skew_sinkoft', False) and getattr(p, '_is_oft', False)
 
         vt_row = None
         vt_col = None
@@ -249,9 +251,12 @@ class SinkSGD_adv(torch.optim.Optimizer):
         grad = _orthogonalize_gradient(p, grad, group["orthogonal_gradient"])
 
         if normed_mt:
+            # Sinkhorn iterative normalization
             if not is_vector:
-                # Sinkhorn iterative normalization
-                grad = apply_sr_sinkhorn(grad, iters=sinkhorn_iterations, p=p, ortho_project=orthogonal_sinkhorn)
+                if not oft_sym:
+                    grad = apply_sr_sinkhorn(grad, iters=sinkhorn_iterations, p=p, ortho_project=orthogonal_sinkhorn)
+                else:
+                    grad = apply_oft_sinkhorn(grad, iters=sinkhorn_iterations, p=p, ortho_project=orthogonal_sinkhorn)
             else:
                 # For vectors, apply sign operation
                 grad = grad.sign_()
@@ -343,9 +348,15 @@ class SinkSGD_adv(torch.optim.Optimizer):
                 update.atan2_(denom)
 
         if not group.get('normed_momentum', False):
+            # Sinkhorn iterative normalization
             if not is_vector:
-                # Sinkhorn iterative normalization
-                update = apply_sr_sinkhorn(update, iters=sinkhorn_iterations, p=p, ortho_project=orthogonal_sinkhorn)
+                if not oft_sym:
+                    update = apply_sr_sinkhorn(update, iters=sinkhorn_iterations, p=p, ortho_project=orthogonal_sinkhorn)
+                else:
+                    if group.get('spectral_normalization', False):
+                        update = apply_spectral_oft_sinkhorn(p, update, iters=sinkhorn_iterations, lr=step_size, state=state, ortho_project=orthogonal_sinkhorn)
+                    else:
+                        update = apply_oft_sinkhorn(update, iters=sinkhorn_iterations, p=p, ortho_project=orthogonal_sinkhorn)
             else:
                 # For vectors, apply sign operation
                 update = update.sign_()
@@ -362,9 +373,13 @@ class SinkSGD_adv(torch.optim.Optimizer):
                 del anchor
 
         update_scaling = step_size
-        if group.get('spectral_normalization', False):
+
+        # If spectral normalization was already applied natively in the OFT sinkhorn, skip the scaling block
+        spectral_sinkhorn_applied = oft_sym and not normed_mt and group.get('spectral_normalization', False)
+
+        if group.get('spectral_normalization', False) and not spectral_sinkhorn_applied:
             update = scale_update(p, update, update_scaling, state=state)
-        else:
+        elif not spectral_sinkhorn_applied:
             if snr_cond:
                 update_scaling = update_scaling * (4/math.pi)
             update.mul_(update_scaling)

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/util/param_update.py

@@ -6,7 +6,7 @@ import torch.nn.functional as F
 
 from typing import Dict, Any
 
-from .scaled_optm import adjust_wds, scale_wd
+from .scaled_optm import adjust_wds
 from .centered_decay import dequantize_anchor
 
 _generators: Dict[torch.device, torch.Generator] = {}
@@ -18,8 +18,8 @@ def _apply_weight_decay(
     p: Tensor,
     state: Dict[str, Any],
     group: Dict[str, Any],
-    eff_wd: float | Tensor | None,
-    eff_cwd: float | Tensor | None,
+    scaled_wd: float | Tensor | None,
+    scaled_cwd: float | Tensor | None,
     wd_target: Tensor | None = None,
     cwd_target: Tensor | None = None,
 ) -> None:
@@ -29,26 +29,26 @@ def _apply_weight_decay(
     cautious = group.get('cautious_wd', False)
 
     # Standard Weight Decay (pulls toward zero)
-    if eff_wd is not None:
+    if scaled_wd is not None:
         if wd_target is None:
             wd_target = p_calc
         # Cautious Weight Decay: only decay if the update pushes in the same direction as the decay
         if cautious:
             mask = (update_calc * p_calc >= 0).to(p_calc.dtype)
-            if isinstance(eff_wd, Tensor):
-                p_calc.addcmul_(wd_target, mask * eff_wd, value=-1.0)
+            if isinstance(scaled_wd, Tensor):
+                p_calc.addcmul_(wd_target, mask * scaled_wd, value=-1.0)
             else:
-                p_calc.addcmul_(wd_target, mask, value=-eff_wd)
+                p_calc.addcmul_(wd_target, mask, value=-scaled_wd)
             del mask
         else:
             # Standard decoupled weight decay
-            if isinstance(eff_wd, Tensor):
-                p_calc.addcmul_(wd_target, eff_wd, value=-1.0)
+            if isinstance(scaled_wd, Tensor):
+                p_calc.addcmul_(wd_target, scaled_wd, value=-1.0)
             else:
-                p_calc.add_(wd_target, alpha=-eff_wd)
+                p_calc.add_(wd_target, alpha=-scaled_wd)
 
     # Centered Weight Decay (pulls toward anchor)
-    if eff_cwd is not None and 'anchor_data' in state:
+    if scaled_cwd is not None and 'anchor_data' in state:
         if cwd_target is not None:
             decay_target = cwd_target
         else:
@@ -59,17 +59,17 @@ def _apply_weight_decay(
         if cautious:
             # Cautious Weight Decay: only decay if the update pushes in the same direction as the decay
             mask = (update_calc * decay_target >= 0).to(p_calc.dtype)
-            if isinstance(eff_cwd, Tensor):
-                p_calc.addcmul_(decay_target, mask * eff_cwd, value=-1.0)
+            if isinstance(scaled_cwd, Tensor):
+                p_calc.addcmul_(decay_target, mask * scaled_cwd, value=-1.0)
             else:
-                p_calc.addcmul_(decay_target, mask, value=-eff_cwd)
+                p_calc.addcmul_(decay_target, mask, value=-scaled_cwd)
             del mask
         else:
             # Standard decoupled weight decay
-            if isinstance(eff_cwd, Tensor):
-                p_calc.addcmul_(decay_target, eff_cwd, value=-1.0)
+            if isinstance(scaled_cwd, Tensor):
+                p_calc.addcmul_(decay_target, scaled_cwd, value=-1.0)
             else:
-                p_calc.add_(decay_target, alpha=-eff_cwd)
+                p_calc.add_(decay_target, alpha=-scaled_cwd)
 
         if cwd_target is None:
             del decay_target
@@ -105,24 +105,18 @@ def apply_parameter_update(
     wd = group["weight_decay"] if wd is None else wd
     cwd = group.get("centered_wd", 0.0)
     wd, cwd = adjust_wds(wd, cwd, p)
-    scaled_wd = group.get("scaled_wd", False)
-    decoupled = scaled_wd or decoupled
 
     # Calculate global decay factor for decoupled vs standard
     decay_factor = (lr / self._init_lr) if decoupled else lr
 
-    eff_wd = (wd * decay_factor) if wd != 0 else None
-    eff_cwd = (cwd * decay_factor) if cwd != 0 else None
-
-    if scaled_wd:
-        eff_wd = scale_wd(eff_wd, p, skip_vectors=True)
-        eff_cwd = scale_wd(eff_cwd, p, skip_vectors=False)
+    scaled_wd = (wd * decay_factor) if wd != 0 else None
+    scaled_cwd = (cwd * decay_factor) if cwd != 0 else None
 
     if wd_scaler is not None:
-        if eff_wd is not None:
-            eff_wd = eff_wd * wd_scaler
-        if eff_cwd is not None:
-            eff_cwd = eff_cwd * wd_scaler
+        if scaled_wd is not None:
+            scaled_wd = scaled_wd * wd_scaler
+        if scaled_cwd is not None:
+            scaled_cwd = scaled_cwd * wd_scaler
 
     state = self.state[p]
 
@@ -135,8 +129,8 @@ def apply_parameter_update(
         cwd_t = cwd_target.float() if cwd_target is not None else None
 
         # Apply weight decay if needed
-        if eff_wd is not None or eff_cwd is not None:
-            _apply_weight_decay(p_fp32, update_fp32, p, state, group, eff_wd, eff_cwd, wd_t, cwd_t)
+        if scaled_wd is not None or scaled_cwd is not None:
+            _apply_weight_decay(p_fp32, update_fp32, p, state, group, scaled_wd, scaled_cwd, wd_t, cwd_t)
 
         # Apply main update
         p_fp32.add_(-update_fp32)
@@ -153,8 +147,8 @@ def apply_parameter_update(
 
     else:
         # Standard path for non-bfloat16 or without stochastic rounding
-        if eff_wd is not None or eff_cwd is not None:
-            _apply_weight_decay(p, update, p, state, group, eff_wd, eff_cwd, wd_target, cwd_target)
+        if scaled_wd is not None or scaled_cwd is not None:
+            _apply_weight_decay(p, update, p, state, group, scaled_wd, scaled_cwd, wd_target, cwd_target)
 
         # Apply main update
         p.add_(-update)

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm/util/scaled_optm.py

@@ -97,25 +97,6 @@ def adjust_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]:
         # Centered WD safely regularizes the delta without collapsing base feature variance.
         return wd, cwd
 
-def scale_wd(wd: float | None, p: torch.Tensor, skip_vectors: bool = False) -> float:
-    """
-    Scale-invariant, dimension-scaled weight decay.
-    """
-    if wd is None:
-        return wd
-    if getattr(p, '_is_oft', False):
-        n_el = p.shape[-1]
-        b = (1.0 + math.sqrt(1.0 + 8.0 * n_el)) / 2.0
-        wd = (2 * wd) / (b - 1)
-        return wd
-
-    is_vector = p.ndim < 2 or getattr(p, '_is_dora_scale', False) or getattr(p, 'is_vector', False) 
-    if is_vector:
-        return 0.0 if skip_vectors else wd
-
-    if p.ndim >= 2:
-        width = p.numel() // p.shape[0]
-        return wd / width
 
 def is_spectral(p: torch.Tensor) -> bool:
     """Determines if a parameter should undergo spectral normalization updates."""
@@ -133,7 +114,7 @@ def init_spectral_norm(state: dict, p: torch.Tensor):
     if getattr(p, '_is_oft', False):
         n_el = p.shape[-1]
         b = int((1.0 + math.sqrt(1.0 + 8.0 * n_el)) / 2.0)
-        _, _, _ = get_cached_structural_tensors(b, p.dtype, p.device)
+        get_cached_structural_tensors(b, p.device)
         gen = param_update.get_generator(p.device)
         batch_size = p.numel() // n_el
         # Initialize v (Right singular vector)

adv_optm-2.6.2.dev2/adv_optm/util/sinkhorn.py

@@ -0,0 +1,275 @@
+import math
+import torch
+
+from . import scaled_optm
+
+def apply_sr_sinkhorn(update: torch.Tensor, iters: int = 5, p: torch.Tensor | None = None, ortho_project: bool = False) -> 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
+        # Stability floor: equivalent to a single-element vector norm lower bound (lb)
+        norm1_lb = 1 / math.sqrt(update_2d.shape[dim])
+        norm1 = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(norm1_lb)
+        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_lb = 1 / math.sqrt(update_2d.shape[1-dim])
+        norm2 = update_2d.norm(p=2, dim=1-dim, keepdim=True).clamp_min_(norm2_lb)
+        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
+    norm_lb = 1 / math.sqrt(update_2d.shape[dim])
+    g_orth_norm = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(norm_lb)
+    scale_factor = target_norm / g_orth_norm
+    return update_2d.mul_(scale_factor)
+
+def get_sinkhorn_wd_scaler(
+    p: torch.Tensor, 
+    row_denom: torch.Tensor | None = None, 
+    col_denom: torch.Tensor | None = None
+):
+    """
+    Computes a structural weight decay multiplier.
+    Penalizes parameters belonging to dominant rows/columns more heavily, 
+    while protecting parameters in under-utilized/noisy rows/columns from decay.
+    """
+    if p.ndim < 2:
+        return 1.0 
+
+    p_2d = p.view(p.shape[0], -1)
+
+    # Lower bounds based on the effective 2D shapes
+    row_lb = 1 / math.sqrt(p_2d.shape[1])
+    col_lb = 1 / math.sqrt(p_2d.shape[0])
+
+    # Get the norms
+    row_norms = torch.linalg.vector_norm(p_2d, ord=2, dim=1, keepdim=True).clamp_min_(row_lb)
+    col_norms = torch.linalg.vector_norm(p_2d, ord=2, dim=0, keepdim=True).clamp_min_(col_lb)
+
+    # Compute the structural scaler
+    row_factor = row_norms.sqrt_()
+    col_factor = col_norms.sqrt_()
+
+    if row_denom is not None and col_denom is not None:
+        # Reshape denominators to ensure safe in-place broadcasting
+        row_denom = row_denom.sqrt().view(p_2d.shape[0], 1)
+        col_denom = col_denom.sqrt().view(1, p_2d.shape[1])
+
+        # High denom (noise) -> smaller angle (protects weights)
+        # Low denom (confident) -> larger angle (decays weights)
+        row_factor.atan2_(row_denom)
+        col_factor.atan2_(col_denom)
+
+    # Outer product: merges the row and column confidences into a 2D matrix
+    wd_scaler = row_factor * col_factor
+
+    # Normalize the scaler so its mean is exactly 1.0
+    wd_scaler.div_(wd_scaler.mean().clamp_min_(1e-12))
+
+    return wd_scaler.view_as(p)
+
+def apply_oft_sinkhorn(
+    update: torch.Tensor, 
+    iters: int = 5, 
+    p: torch.Tensor | None = None, 
+    ortho_project: bool = False
+):
+    n_el = update.shape[-1]
+    block_size = int((1 + math.sqrt(1 + 8 * n_el)) / 2)
+    device, dtype = update.device, update.dtype
+    rows, cols = scaled_optm.get_cached_structural_tensors(block_size, device)
+    batch_size = update.shape[0]
+
+    # Initialize matrices
+    G = torch.zeros(batch_size, block_size, block_size, device=device, dtype=dtype)
+    batch_idx = torch.arange(batch_size, device=device)[:, None]
+
+    # Construct skew-symmetric gradient matrix G
+    G = G.index_put((batch_idx, rows, cols), update)
+    G = G - G.transpose(-2, -1)
+
+    # If OrthoGrad is enabled, construct the skew-symmetric parameter matrix Q
+    if ortho_project and p is not None:
+        Q = torch.zeros_like(G)
+        p_flat = p.view(batch_size, -1)
+        Q = Q.index_put((batch_idx, rows, cols), p_flat)
+        Q = Q - Q.transpose(-2, -1)
+        # Compute the squared Frobenius norm for each block (batch-wise)
+        Q_norm_sq = torch.sum(Q * Q, dim=(1, 2), keepdim=True).add_(1e-30)
+        # Theoretical norm of a Sinkhorn-normalized matrix to restore magnitude
+        target_norm = math.sqrt(block_size) 
+
+    for _ in range(iters):
+        norms = torch.linalg.vector_norm(G, ord=2, dim=2, keepdim=True).clamp_min_(1e-12)
+        S = norms.sqrt_()
+        G.div_(S * S.mT)
+
+        # OrthoGrad step
+        if ortho_project and p is not None:
+            # Batch-wise Frobenius inner product
+            dot_prod = torch.sum(Q * G, dim=(1, 2), keepdim=True)
+            proj = dot_prod / Q_norm_sq
+            # Subtract projection
+            G.addcmul_(proj, Q, value=-1.0)
+            # Restore magnitude
+            g_orth_norm = torch.linalg.vector_norm(G, ord=2, dim=(1, 2), keepdim=True).clamp_min_(1e-12)
+            G.mul_(target_norm / g_orth_norm)
+
+    update.copy_(G[batch_idx, rows, cols])
+
+    # Global normalization
+    norm = torch.linalg.vector_norm(update).clamp_min_(1e-12)
+    target_norm_global = math.sqrt(update.numel())
+    update.mul_(target_norm_global / norm)
+
+    return update
+
+@torch.no_grad()
+def apply_spectral_oft_sinkhorn(
+    p: torch.Tensor,
+    update: torch.Tensor,
+    lr: float,
+    state: dict,
+    iters: int = 5,
+    ortho_project: bool = True
+) -> torch.Tensor:
+    """
+    Applies Spectral Normalization directly on the skew-symmetric gradient.
+    """
+    n_el = p.shape[-1]
+    block_size = int((1 + math.sqrt(1 + 8 * n_el)) / 2)
+    device, dtype = p.device, p.dtype
+    rows, cols = scaled_optm.get_cached_structural_tensors(block_size, device)
+
+    orig_shape = p.shape
+
+    # Align the scale of p with the forward pass
+    scale_factor = getattr(p, '_oft_scale_factor', 1.0)
+    batch_size = update.shape[0]
+
+    # Initialize matrices
+    G = torch.zeros(batch_size, block_size, block_size, device=device, dtype=dtype)
+    batch_idx = torch.arange(batch_size, device=device)[:, None]
+
+    # Construct skew-symmetric gradient matrix G
+    G = G.index_put((batch_idx, rows, cols), update)
+    G = G - G.transpose(-2, -1)
+
+    # If OrthoGrad is enabled, construct the skew-symmetric parameter matrix Q
+    if ortho_project and p is not None:
+        Q = torch.zeros_like(G)
+        p_flat = p.view(batch_size, -1)
+        Q = Q.index_put((batch_idx, rows, cols), p_flat)
+        Q = Q - Q.transpose(-2, -1)
+        # Compute the squared Frobenius norm for each block (batch-wise)
+        Q_norm_sq = torch.sum(Q * Q, dim=(1, 2), keepdim=True).add_(1e-30)
+        # Theoretical norm of a Sinkhorn-normalized matrix to restore magnitude
+        target_norm = math.sqrt(block_size) 
+
+    for _ in range(iters):
+        norms = torch.linalg.vector_norm(G, ord=2, dim=2, keepdim=True).clamp_min_(1e-12)
+        S = norms.sqrt_()
+        G.div_(S * S.mT)
+
+        # OrthoGrad step
+        if ortho_project and p is not None:
+            # Batch-wise Frobenius inner product
+            dot_prod = torch.sum(Q * G, dim=(1, 2), keepdim=True)
+            proj = dot_prod / Q_norm_sq
+            # Subtract projection
+            G.addcmul_(proj, Q, value=-1.0)
+            # Restore magnitude
+            g_orth_norm = torch.linalg.vector_norm(G, ord=2, dim=(1, 2), keepdim=True).clamp_min_(1e-12)
+            G.mul_(target_norm / g_orth_norm)
+
+    update.copy_(G[batch_idx, rows, cols])
+
+    # Spectral Normalization on G
+    u_state = state['spectral_u'].unsqueeze(-1).to(dtype)
+    v_state = state['spectral_v'].unsqueeze(-1).to(dtype)
+    # Power Iteration step to estimate the largest singular value (sigma)
+    # Update v (Right Singular Vector)
+    v_raw = torch.bmm(G.mT, u_state)
+    v_norm = torch.linalg.vector_norm(v_raw, dim=1, keepdim=True)
+    candidate_v = v_raw / v_norm.clamp_min(1e-8)
+    next_v = torch.where(v_norm >= 1e-6, candidate_v, v_state)
+    # Update u (Left Singular Vector)
+    u_raw = torch.bmm(G, next_v)
+    u_norm = torch.linalg.vector_norm(u_raw, dim=1, keepdim=True)
+    candidate_u = u_raw / u_norm.clamp_min(1e-8)
+    next_u = torch.where(u_norm >= 1e-6, candidate_u, u_state)
+    state['spectral_v'].copy_(next_v.squeeze(-1))
+    state['spectral_u'].copy_(next_u.squeeze(-1))
+
+    # Estimate sigma (The spectral norm) for each block
+    sigma = torch.sum(next_u * u_raw, dim=1, keepdim=True)
+
+    # Squeeze out the last dimension so shape becomes (batch_size, 1)
+    sigma = sigma.squeeze(-1) 
+
+    target_scale = 0.5 * scale_factor
+    spectral_eps = 1.0 / (2.0 * math.sqrt(block_size))
+
+    # Apply the clamp and scaling block-wise
+    scale = lr * (target_scale / sigma.clamp_min(spectral_eps))
+
+    return update.mul_(scale).view(orig_shape)

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-2.6.1.dev3 → adv_optm-2.6.2.dev2}/setup.py

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

adv_optm-2.6.1.dev3/adv_optm/util/sinkhorn.py

@@ -1,126 +0,0 @@
-import math
-import torch
-
-def apply_sr_sinkhorn(update: torch.Tensor, iters: int = 5, p: torch.Tensor | None = None, ortho_project: bool = False) -> 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
-        # Stability floor: equivalent to a single-element vector norm lower bound (lb)
-        norm1_lb = 1 / math.sqrt(update_2d.shape[dim])
-        norm1 = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(norm1_lb)
-        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_lb = 1 / math.sqrt(update_2d.shape[1-dim])
-        norm2 = update_2d.norm(p=2, dim=1-dim, keepdim=True).clamp_min_(norm2_lb)
-        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
-    norm_lb = 1 / math.sqrt(update_2d.shape[dim])
-    g_orth_norm = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(norm_lb)
-    scale_factor = target_norm / g_orth_norm
-    return update_2d.mul_(scale_factor)
-
-def get_sinkhorn_wd_scaler(
-    p: torch.Tensor, 
-    row_denom: torch.Tensor | None = None, 
-    col_denom: torch.Tensor | None = None
-):
-    """
-    Computes a structural weight decay multiplier.
-    Penalizes parameters belonging to dominant rows/columns more heavily, 
-    while protecting parameters in under-utilized/noisy rows/columns from decay.
-    """
-    if p.ndim < 2:
-        return 1.0 
-
-    p_2d = p.view(p.shape[0], -1)
-
-    # Lower bounds based on the effective 2D shapes
-    row_lb = 1 / math.sqrt(p_2d.shape[1])
-    col_lb = 1 / math.sqrt(p_2d.shape[0])
-
-    # Get the norms
-    row_norms = torch.linalg.vector_norm(p_2d, ord=2, dim=1, keepdim=True).clamp_min_(row_lb)
-    col_norms = torch.linalg.vector_norm(p_2d, ord=2, dim=0, keepdim=True).clamp_min_(col_lb)
-
-    # Compute the structural scaler
-    row_factor = row_norms.sqrt_()
-    col_factor = col_norms.sqrt_()
-
-    if row_denom is not None and col_denom is not None:
-        # Reshape denominators to ensure safe in-place broadcasting
-        row_denom = row_denom.sqrt().view(p_2d.shape[0], 1)
-        col_denom = col_denom.sqrt().view(1, p_2d.shape[1])
-
-        # High denom (noise) -> smaller angle (protects weights)
-        # Low denom (confident) -> larger angle (decays weights)
-        row_factor.atan2_(row_denom)
-        col_factor.atan2_(col_denom)
-
-    # Outer product: merges the row and column confidences into a 2D matrix
-    wd_scaler = row_factor * col_factor
-
-    # Normalize the scaler so its mean is exactly 1.0
-    wd_scaler.div_(wd_scaler.mean().clamp_min_(1e-12))
-
-    return wd_scaler.view_as(p)