adv-optm 2.4.dev24__tar.gz → 2.5__tar.gz

{adv_optm-2.4.dev24 → adv_optm-2.5}/PKG-INFO

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

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "SinkSGD_adv",
 ]
 
-__version__ = "2.4.dev24"
+__version__ = "2.5"

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/optim/AdaMuon_adv.py

@@ -57,7 +57,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
             (default: (3.4445, -4.7750, 2.0315)).
         stochastic_rounding (bool): whether to use stochastic rounding for
             BF16 parameter updates (default: True).
-        orthogonal_gradient (bool): whether to use OrthoGrad.  (default: False)
+        orthogonal_gradient (str): whether to use OrthoGrad variants. 'disabled': off.
+        'flattened': Standard vectorized OrthoGrad. 'iterative': Matrix-wise rank-2 OrthoGrad. (default: disabled)
         nesterov (bool): enables Nesterov momentum (default: False).
         use_atan2 (bool): whether to use the atan2 update rule. (default: False)
         vector_reshape (bool): whether to reshape 1D vectors into 2D
@@ -114,7 +115,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         adam_fisher_wd (bool): Fisher Adam (FAdam) weight decay for the AdamW part. (default: False)
         adam_use_bias_correction (bool): Bias correction for AdamW.
         adam_use_atan2 (bool): Atan2 update rule for AdamW.
-        adam_orthogonal_gradient (bool): OrthoGrad for AdamW.
+        adam_orthogonal_gradient (str): OrthoGrad for AdamW.
         adam_nesterov (bool): Nesterov momentum for AdamW. (default: False)
         adam_nesterov_coef (float, optional): Nesterov coefficient for AdamW. (default: None)
         adam_kourkoutas_beta (bool): Kourkoutas-β for AdamW.
@@ -149,7 +150,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
-        orthogonal_gradient: bool = False,
+        orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         # Adam_atan2 (scale invariant)
         use_atan2: bool = False,
         # NorMuon
@@ -190,7 +191,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         adam_fisher_wd: bool = False,
         adam_use_bias_correction: bool = True,
         adam_use_atan2: bool = False,
-        adam_orthogonal_gradient: bool = False,
+        adam_orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         adam_nesterov: bool = False,
         adam_nesterov_coef: float | None = None,
         adam_kourkoutas_beta: bool = False,
@@ -213,7 +214,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             print("Warning: spectral_normalization is incompatible with rms_rescaling, Disabling rms_rescaling.")
             rms_rescaling = False
         if spectral_normalization and accelerated_ns:
-            ValueError("spectral_normalization violates accelerated Newton-Schulz assumptions. Pick one of them.")
+            raise ValueError("spectral_normalization violates accelerated Newton-Schulz assumptions. Pick one of them.")
 
         # Legacy backwards compatibility support for `nnmf_factor=True`
         if nnmf_factor:
@@ -515,8 +516,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             grad = approx_mars(grad, state['last_grad'], group['mars_gamma'], beta1)
 
 
-        if group.get("orthogonal_gradient"):
-            grad = _orthogonalize_gradient(p, grad)
+        grad = _orthogonalize_gradient(p, grad, group.get("orthogonal_gradient"))
 
         if state['factored']: # Factored Muon
             d1, d2 = state['effective_shape']

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/optim/AdamW_adv.py

@@ -45,7 +45,8 @@ class AdamW_adv(torch.optim.Optimizer):
         stochastic_rounding (bool): whether to use stochastic
             rounding for BF16 parameter updates (default: True).
         use_atan2 (bool): whether to use the atan2 update rule. (default: False)
-        orthogonal_gradient (bool): whether to use OrthoGrad.  (default: False)
+        orthogonal_gradient (str): whether to use OrthoGrad variants. 'disabled': off.
+        'flattened': Standard vectorized OrthoGrad. 'iterative': Matrix-wise rank-2 OrthoGrad. (default: disabled)
         normed_momentum (bool): whether to compute the first moment on the normalized gradient. (default: False)
         kourkoutas_beta (bool): whether to enable the layer-wise dynamic β₂ logic.
             If `False`, the optimizer behaves as standard AdamW. (default: False)
@@ -104,7 +105,7 @@ class AdamW_adv(torch.optim.Optimizer):
         # Adam_atan2 (scale invariant)
         use_atan2: bool = False,
         # OrthoGrad
-        orthogonal_gradient: bool = False,
+        orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         # Nesterov momentum
         nesterov: bool = False,
         nesterov_coef: float | None = None,
@@ -326,8 +327,7 @@ class AdamW_adv(torch.optim.Optimizer):
     def _step_parameter(self, p, grad, state, group, step_size, beta1, beta2, sqrt_bias_correction2, random_int_tensor, random_int_state_tensor):
         grad = upcast_grad_for_precision(grad, state, group['state_precision'])
 
-        if group["orthogonal_gradient"]:
-            grad = _orthogonalize_gradient(p, grad)
+        grad = _orthogonalize_gradient(p, grad, group["orthogonal_gradient"])
 
         nesterov = group.get('nesterov', False)
         nesterov_coef = group.get('nesterov_coef', None)
@@ -462,7 +462,7 @@ class AdamW_adv(torch.optim.Optimizer):
         else:
             update.mul_(update_scaling)
 
-        param_update.apply_parameter_update(self, p, group, update, step_size, random_int_tensor=random_int_tensor, wd_scaler=wd_scaler)
+        param_update.apply_parameter_update(self, p, group, update, group['lr'], random_int_tensor=random_int_tensor, wd_scaler=wd_scaler)
 
     def compile(self, *args, **kwargs):
         self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/optim/Adopt_adv.py

@@ -108,7 +108,7 @@ class Adopt_adv(torch.optim.Optimizer):
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
-        orthogonal_gradient: bool = False,
+        orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         # Nesterov momentum
         nesterov: bool = False,
         nesterov_coef: float | None = None,
@@ -158,7 +158,7 @@ class Adopt_adv(torch.optim.Optimizer):
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
-            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
+            "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,
             "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
@@ -172,7 +172,6 @@ class Adopt_adv(torch.optim.Optimizer):
         self.clip_lambda = clip_lambda
         self.stochastic_rounding = stochastic_rounding
         self.use_atan2 = use_atan2
-        self.orthogonal_gradient = orthogonal_gradient
         self.kourkoutas_beta = kourkoutas_beta
         self.layer_key_fn = layer_key_fn
         self._init_lr = lr if lr > 0 else 1
@@ -237,7 +236,7 @@ class Adopt_adv(torch.optim.Optimizer):
             dtype = torch.float32 if (state['factored'] or req_precision == 'factored') else p.dtype
 
             vt_dtype = torch.float32 if (state['factored'] or state['factored_2nd'] or req_precision in ['factored', 'bf16_sr', 'int8_sr']) else dtype
-            vt_init = grad.pow(2).to(vt_dtype) * (1 - group['betas'][1])
+            vt_init = grad.pow(2).to(vt_dtype)
 
             if state['factored']:
                 state['effective_shape'] = _get_effective_shape(p.numel())
@@ -329,8 +328,7 @@ class Adopt_adv(torch.optim.Optimizer):
     def _step_parameter(self, p, grad, state, group, lr, beta1, beta2, random_int_tensor, random_int_state_tensor):
         grad = upcast_grad_for_precision(grad, state, group['state_precision'])
 
-        if self.orthogonal_gradient:
-            grad = _orthogonalize_gradient(p, grad)
+        grad = _orthogonalize_gradient(p, grad, group["orthogonal_gradient"])
 
         nesterov = group.get('nesterov', False)
         nesterov_coef = group.get('nesterov_coef', None)

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/optim/Lion_adv.py

@@ -67,7 +67,7 @@ class Lion_adv(torch.optim.Optimizer):
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
-        orthogonal_gradient: bool = False,
+        orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         # Lion-k
         kappa_p: float = 1.0,
         auto_kappa_p: bool = False,
@@ -213,8 +213,9 @@ class Lion_adv(torch.optim.Optimizer):
     def _step_parameter(self, p, grad, state, group, lr, random_int_tensor, random_noise_tensor):
         if grad.dtype != torch.float32 and state['factored']:
             grad = grad.float()
-        if group["orthogonal_gradient"]:
-            grad = _orthogonalize_gradient(p, grad)
+        is_vector = p.ndim < 2 or getattr(p, '_is_dora_scale', False) or getattr(p, 'is_vector', False)
+
+        grad = _orthogonalize_gradient(p, grad, group["orthogonal_gradient"])
 
         # Lion-K Logic
         kappa_p = group.get("kappa_p", 1.0)
@@ -250,7 +251,7 @@ class Lion_adv(torch.optim.Optimizer):
             update = update.view(p.shape)
 
             if group.get('stochastic_sign', False):
-                update = apply_stochastic_sign_(update, noise=random_noise_tensor)
+                update = apply_stochastic_sign_(update, noise=random_noise_tensor, is_vector=is_vector)
             else:
                 update = _get_lion_k_update(update, kappa_p)
 
@@ -265,7 +266,7 @@ class Lion_adv(torch.optim.Optimizer):
             exp_avg.lerp_(grad, 1 - beta2)
 
             if group.get('stochastic_sign', False):
-                update = apply_stochastic_sign_(update, noise=random_noise_tensor)
+                update = apply_stochastic_sign_(update, noise=random_noise_tensor, is_vector=is_vector)
             else:
                 update = _get_lion_k_update(update, kappa_p)
 

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/optim/Muon_adv.py

@@ -39,7 +39,8 @@ class Muon_adv(torch.optim.Optimizer):
             (default: (3.4445, -4.7750, 2.0315)).
         stochastic_rounding (bool): whether to use stochastic rounding for
             BF16 parameter updates (default: True).
-        orthogonal_gradient (bool): whether to use OrthoGrad.  (default: False)
+        orthogonal_gradient (str): whether to use OrthoGrad variants. 'disabled': off.
+        'flattened': Standard vectorized OrthoGrad. 'iterative': Matrix-wise rank-2 OrthoGrad. (default: disabled)
         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
@@ -89,7 +90,7 @@ class Muon_adv(torch.optim.Optimizer):
         adam_fisher_wd (bool): Fisher Adam (FAdam) weight decay for the AdamW part. (default: False)
         adam_use_bias_correction (bool): Bias correction for AdamW.
         adam_use_atan2 (bool): Atan2 update rule for AdamW.
-        adam_orthogonal_gradient (bool): OrthoGrad for AdamW.
+        adam_orthogonal_gradient (str): OrthoGrad for AdamW.
         adam_nesterov (bool): Nesterov momentum for AdamW. (default: False)
         adam_nesterov_coef (float, optional): Nesterov coefficient for AdamW. (default: None)
         adam_kourkoutas_beta (bool): Kourkoutas-β for AdamW.
@@ -121,7 +122,7 @@ class Muon_adv(torch.optim.Optimizer):
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
-        orthogonal_gradient: bool = False,
+        orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         # RMS Rescaling
         rms_rescaling: bool = True,
         # SMMF factorization
@@ -159,7 +160,7 @@ class Muon_adv(torch.optim.Optimizer):
         adam_fisher_wd: bool = False,
         adam_use_bias_correction: bool = True,
         adam_use_atan2: bool = False,
-        adam_orthogonal_gradient: bool = False,
+        adam_orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         adam_nesterov: bool = False,
         adam_nesterov_coef: float | None = None,
         adam_kourkoutas_beta: bool = False,
@@ -186,7 +187,7 @@ class Muon_adv(torch.optim.Optimizer):
             print("Warning: spectral_normalization is incompatible with rms_rescaling, Disabling rms_rescaling.")
             rms_rescaling = False
         if spectral_normalization and accelerated_ns:
-            ValueError("spectral_normalization violates accelerated Newton-Schulz assumptions. Pick one of them.")
+            raise ValueError("spectral_normalization violates accelerated Newton-Schulz assumptions. Pick one of them.")
 
         # Legacy backwards compatibility support for `nnmf_factor=True`
         if nnmf_factor:
@@ -457,8 +458,7 @@ class Muon_adv(torch.optim.Optimizer):
         if grad.dtype != torch.float32 and state.get('factored', False):
             grad = grad.float()
 
-        if group.get("orthogonal_gradient"):
-            grad = _orthogonalize_gradient(p, grad)
+        grad = _orthogonalize_gradient(p, grad, group.get("orthogonal_gradient"))
 
         if state['factored']: # Factored Muon
             d1, d2 = state['effective_shape']

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/optim/Prodigy_adv.py

@@ -43,7 +43,8 @@ class Prodigy_adv(torch.optim.Optimizer):
         stochastic_rounding (bool): whether to use stochastic
             rounding for BF16 parameter updates (default: True).
         use_atan2 (bool): whether to use the atan2 update rule. (default: False)
-        orthogonal_gradient (bool): whether to use OrthoGrad.  (default: False)
+        orthogonal_gradient (str): whether to use OrthoGrad variants. 'disabled': off.
+        'flattened': Standard vectorized OrthoGrad. 'iterative': Matrix-wise rank-2 OrthoGrad. (default: disabled)
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
         factored_2nd (bool): whether to keep the first moment uncompressed (dense)
@@ -119,7 +120,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         # Adam_atan2 (scale invariant)
         use_atan2: bool = False,
         # OrthoGrad
-        orthogonal_gradient: bool = False,
+        orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         # Nesterov momentum
         nesterov: bool = False,
         nesterov_coef: float | None = None,
@@ -371,8 +372,7 @@ class Prodigy_adv(torch.optim.Optimizer):
     def _step_parameter(self, p, grad, state, group, beta2, d, dlr, random_int_tensor, random_int_state_tensor):
         grad = upcast_grad_for_precision(grad, state, group['state_precision'])
 
-        if group["orthogonal_gradient"]:
-            grad = _orthogonalize_gradient(p, grad)
+        grad = _orthogonalize_gradient(p, grad, group["orthogonal_gradient"])
 
         nesterov = group.get('nesterov', False)
         nesterov_coef = group.get('nesterov_coef', None)

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/optim/SignSGD_adv.py

@@ -62,7 +62,7 @@ class SignSGD_adv(torch.optim.Optimizer):
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
-        orthogonal_gradient: bool = False,
+        orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         # Stochastic Sign Operator
         stochastic_sign: bool = False,
         # Nesterov momentum
@@ -171,7 +171,7 @@ class SignSGD_adv(torch.optim.Optimizer):
     def __init_state(self, p, group):
         state = self.state[p]
         # State Initialization
-        if group["momentum"] > 0 and len(state) == 0:
+        if 'step' not in state:
             req_precision = group['state_precision']
             is_vector = len(p.shape) == 1 and not group['vector_reshape']
 
@@ -259,8 +259,7 @@ class SignSGD_adv(torch.optim.Optimizer):
         wd_target = None
         cwd_target = None
 
-        if group["orthogonal_gradient"]:
-            grad = _orthogonalize_gradient(p, grad)
+        grad = _orthogonalize_gradient(p, grad, group["orthogonal_gradient"])
 
         if normed_mt:
             if sso:
@@ -280,16 +279,18 @@ class SignSGD_adv(torch.optim.Optimizer):
                 if snr_cond:
                     denom = (1.0 - exp_avg.square()).clamp_min_(1e-30).sqrt_().view_as(p)
 
+                if nesterov and normed_mt:
+                    # Scale the normalized gradient using empirical buffer magnitude (SNR recovery)
+                    normed_grad = exp_avg.abs().mul_(grad_reshaped)
+
                 exp_avg.lerp_(grad_reshaped, 1 - momentum)
 
                 if nesterov:
                     nv_coef = momentum if nesterov_coef is None else nesterov_coef
                     if normed_mt:
-                        # Scale the normalized gradient down to match the buffer's variance
-                        ema_std = math.sqrt((1 - momentum) / (1 + momentum))
-                        raw_update = (grad_reshaped * ema_std).lerp_(exp_avg, nv_coef)
+                        raw_update = normed_grad.lerp_(exp_avg, nv_coef)
                     else:
-                        raw_update = grad.lerp(exp_avg, nv_coef)
+                        raw_update = grad_reshaped.lerp(exp_avg, nv_coef)
                 else:
                     raw_update = exp_avg.clone()
 
@@ -309,14 +310,16 @@ class SignSGD_adv(torch.optim.Optimizer):
                 if snr_cond:
                     denom = (1.0 - exp_avg.square()).clamp_min_(1e-30).sqrt_()
 
+                if nesterov and normed_mt:
+                    # Scale the normalized gradient using empirical buffer magnitude (SNR recovery)
+                    normed_grad = exp_avg.abs().mul_(grad)
+
                 exp_avg.lerp_(grad, 1 - momentum)
 
                 if nesterov:
                     nv_coef = momentum if nesterov_coef is None else nesterov_coef
                     if normed_mt:
-                        # Scale the normalized gradient down to match the buffer's variance
-                        ema_std = math.sqrt((1 - momentum) / (1 + momentum))
-                        raw_update = (grad * ema_std).lerp_(exp_avg, nv_coef)
+                        raw_update = normed_grad.lerp_(exp_avg, nv_coef)
                     else:
                         raw_update = grad.lerp(exp_avg, nv_coef)
                 else:
@@ -340,7 +343,7 @@ class SignSGD_adv(torch.optim.Optimizer):
         if group.get('geometric_wd', False) and group["weight_decay"] > 0 :
             wd_target = get_signsgd_wd_target(p, denom=denom, stochastic_sign=sso, noise=random_noise_tensor, is_vector=is_vector)
 
-            if group.get('centered_wd', 0.0) > 0 and 'anchor_type' in state:
+            if group.get('centered_wd', 0.0) > 0 and 'anchor_data' in state:
                 anchor = dequantize_anchor(p, state, group, p.dtype)
                 cwd_target = get_signsgd_wd_target(p.sub(anchor), denom=denom, stochastic_sign=sso, noise=random_noise_tensor, is_vector=is_vector)
                 del anchor

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/optim/SinkSGD_adv.py

@@ -69,7 +69,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
-        orthogonal_gradient: bool = False,
+        orthogonal_gradient: str = 'disabled', # 'flattened', 'iterative'
         # Spectral Normed Optimizer
         spectral_normalization: bool = False,
         # Centered WD
@@ -89,8 +89,8 @@ class SinkSGD_adv(torch.optim.Optimizer):
             raise ValueError(f"Momentum should be >= 0.0. Got {momentum}")
         if not (weight_decay >= 0.0):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
-        if snr_cond and not normed_momentum:
-            raise NotImplementedError(f"snr_cond is intended to be used with normed_momentum")
+        if snr_cond and not normed_momentum and not momentum > 0:
+            raise NotImplementedError(f"snr_cond is intended to be used with normed_momentum.")
 
         state_precision = state_precision.lower()
         valid_precisions = {"auto", "fp32", "factored", "bf16_sr", "fp16", "int8_sr"}
@@ -237,8 +237,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
         wd_target = None
         cwd_target = None
 
-        if group["orthogonal_gradient"]:
-            grad = _orthogonalize_gradient(p, grad)
+        grad = _orthogonalize_gradient(p, grad, group["orthogonal_gradient"])
 
         if normed_mt:
             if not is_vector:
@@ -264,6 +263,10 @@ class SinkSGD_adv(torch.optim.Optimizer):
                     else:
                         denom = (1.0 - buf.square()).clamp_min_(1e-30).sqrt_().view_as(p)
 
+                if nesterov and normed_mt:
+                    # Scale the normalized gradient using empirical buffer magnitude (SNR recovery)
+                    normed_grad = buf.abs().mul_(grad_reshaped)
+
                 buf.lerp_(grad_reshaped, 1 - momentum)
 
                 # Factorize updated buffer
@@ -272,9 +275,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
                 if nesterov:
                     nv_coef = momentum if nesterov_coef is None else nesterov_coef
                     if normed_mt:
-                        # Scale the normalized gradient down to match the buffer's variance
-                        ema_std = math.sqrt((1 - momentum) / (1 + momentum))
-                        update = (grad_reshaped * ema_std).lerp_(buf, nv_coef)
+                        update = normed_grad.lerp_(buf, nv_coef)
                     else:
                         update = grad_reshaped.lerp(buf, nv_coef)
                 else:
@@ -299,6 +300,10 @@ class SinkSGD_adv(torch.optim.Optimizer):
                     else:
                         denom = (1.0 - buf.square()).clamp_min_(1e-30).sqrt_()
 
+                if nesterov and normed_mt:
+                    # Scale the normalized gradient using empirical buffer magnitude (SNR recovery)
+                    normed_grad = buf.abs().mul_(grad)
+
                 buf.lerp_(grad, 1 - momentum)
 
                 set_state(state, 'momentum_buffer', buf, actual_precision, random_int_state_tensor)
@@ -306,9 +311,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
                 if nesterov:
                     nv_coef = momentum if nesterov_coef is None else nesterov_coef
                     if normed_mt:
-                        # Scale the normalized gradient down to match the buffer's variance
-                        ema_std = math.sqrt((1 - momentum) / (1 + momentum))
-                        update = (grad * ema_std).lerp_(buf, nv_coef)
+                        update = normed_grad.lerp_(buf, nv_coef)
                     else:
                         update = grad.lerp(buf, nv_coef)
                 else:
@@ -342,7 +345,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
                     wd_scaler = get_sinkhorn_wd_scaler(p, row_denom=vt_row, col_denom=vt_col)
                 else:
                     wd_target = get_signsgd_wd_target(p, denom=denom)
-            if is_vector and group.get('centered_wd', 0.0) > 0 and 'anchor_type' in state:
+            if is_vector and group.get('centered_wd', 0.0) > 0 and 'anchor_data' in state:
                 anchor = dequantize_anchor(p, state, group, p.dtype)
                 cwd_target = get_signsgd_wd_target(p.sub(anchor), denom=denom)
                 del anchor

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/util/Muon_AuxAdam.py

@@ -71,8 +71,7 @@ def _init_auxadam_state(self, p, group):
 def _adam_step_parameter(self, p, grad, state, group, beta1_adam, beta2_adam, sqrt_bias_correction2, step_size, random_int_tensor, random_int_state_tensor=None):
     grad = upcast_grad_for_precision(grad, state, group.get('adam_state_precision', 'auto'))
 
-    if group.get("adam_orthogonal_gradient"):
-        grad = _orthogonalize_gradient(p, grad)
+    grad = _orthogonalize_gradient(p, grad, group.get("adam_orthogonal_gradient"))
 
     if hasattr(self, 'kourkoutas_helper') and self.kourkoutas_helper:
         # Accumulate current grad's norm for the *next* step
@@ -190,4 +189,4 @@ def _adam_step_parameter(self, p, grad, state, group, beta1_adam, beta2_adam, sq
     else:
         update.mul_(update_scaling)
 
-    param_update.apply_parameter_update(self, p, group, update, step_size, group["adam_weight_decay"], random_int_tensor=random_int_tensor, wd_scaler=wd_scaler)
+    param_update.apply_parameter_update(self, p, group, update, group['lr'], group["adam_weight_decay"], random_int_tensor=random_int_tensor, wd_scaler=wd_scaler)

adv_optm-2.5/adv_optm/util/OrthoGrad.py

@@ -0,0 +1,92 @@
+import torch
+import math
+
+def _orthogonalize_gradient(p: torch.Tensor, grad: torch.Tensor, mode: str) -> torch.Tensor:
+    """
+    Projects the gradient `grad` to be orthogonal to the parameter `p`.
+    Supports two modes: 'flattened' (vectorized) and 'iterative' (matrix-wise).
+    """
+    if mode == 'disabled':
+        return grad
+    elif mode == 'flattened':
+        return flattened_ortho_project(p, grad)
+    elif mode == 'iterative':
+        return iterative_ortho_project(p, grad, iters=3)
+
+def flattened_ortho_project(p: torch.Tensor, grad: torch.Tensor) -> torch.Tensor:
+    """
+    Projects the flattened gradient `grad` to be orthogonal to the flattened parameter `p`.
+    Modified from:
+    https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability/blob/720d2444df12b851d6cb417ab08cf125c822b2ae/orthograd.py
+    """
+    original_shape = grad.shape
+    original_dtype = grad.dtype
+    w = p.view(-1).float()
+    g = grad.view(-1).float()
+    w_norm_sq = torch.dot(w, w).add_(1e-30)
+    proj = torch.dot(w, g) / w_norm_sq
+    g_orth = g.sub(w * proj)
+    g_norm = g.norm(2)
+    g_orth_norm = g_orth.norm(2).add_(1e-30)
+    g_orth_scaled = g_orth * (g_norm / g_orth_norm)
+    return g_orth_scaled.view(original_shape).to(original_dtype)
+
+
+def iterative_ortho_project(p: torch.Tensor, grad: torch.Tensor, iters: int = 3) -> torch.Tensor:
+    """
+    Applies iterative alternating orthogonal projection to a 2D matrix.
+    Projects the grad to be orthogonal to the parameter matrix along
+    rows and columns sequentially, alternating dimensions.
+    Inspired from Sinkhorn algorithm, 2-3 iterations is enough to converge
+    to cosine similarity of -1e4 to -1e-6 for every row/col (semi orthogonal).
+    """
+    # 1D Vector Case fallback to the standard OrthoGrad
+    is_vector = p.ndim < 2 or getattr(p, '_is_dora_scale', False) or getattr(p, 'is_vector', False)
+    if is_vector:
+        return _orthogonalize_gradient(p, grad)
+
+    original_shape = grad.shape
+
+    # 2D+ Matrix Case
+    grad_2d = grad.view(grad.shape[0], -1)
+    param_2d = p.view(p.shape[0], -1)
+
+    m, n = grad_2d.shape
+
+    # Dynamically determine the order based on aspect ratio
+    row_first = m > n
+    dim = 0 if row_first else 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)
+
+    for _ in range(iters):
+        # First dimension
+        grad_2d = _ortho_normed_dim(param_2d, grad_2d, p_norm_sq_dim, dim)
+        # Second dimension
+        grad_2d = _ortho_normed_dim(param_2d, grad_2d, p_norm_sq_adim, 1 - dim)
+
+    return grad_2d.view(original_shape)
+
+
+def _ortho_normed_dim(p_2d: torch.Tensor, grad_2d: torch.Tensor, p_norm_sq: torch.Tensor, dim: int) -> torch.Tensor:
+    """
+    Projects the grad to be orthogonal to p along 'dim' and dynamically restores 
+    the original magnitude of that dimension pre-projection.
+    """
+    # Record target magnitude before projection
+    norm_lb = 1 / math.sqrt(grad_2d.shape[dim])
+    target_norm = grad_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(norm_lb)
+
+    # Project: g_orth = g - (p * <p, g> / ||p||^2)
+    dot_prod = torch.sum(p_2d * grad_2d, dim=dim, keepdim=True)
+    proj = dot_prod / p_norm_sq
+
+    # In-place subtraction: grad_2d = grad_2d - (proj * p_2d)
+    # Standard gamma is -1, but -1.01 proved to converge faster
+    grad_2d.addcmul_(proj, p_2d, value=-1.01)
+
+    # Magnitude Preservation
+    g_orth_norm = grad_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(norm_lb)
+    scale_factor = target_norm / g_orth_norm
+    return grad_2d.mul_(scale_factor)

{adv_optm-2.4.dev24 → adv_optm-2.5}/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_wds
+from .scaled_optm import adjust_wds
 from .centered_decay import dequantize_anchor
 
 _generators: Dict[torch.device, torch.Generator] = {}
@@ -48,7 +48,7 @@ def _apply_weight_decay(
                 p_calc.add_(wd_target, alpha=-scaled_wd)
 
     # Centered Weight Decay (pulls toward anchor)
-    if scaled_cwd is not None and 'anchor_type' in state:
+    if scaled_cwd is not None and 'anchor_data' in state:
         if cwd_target is not None:
             decay_target = cwd_target
         else:
@@ -330,7 +330,7 @@ def _copy_int8_sym_blockwise_stochastic_core_(
     target: torch.Tensor,
     source: torch.Tensor,
     scales: torch.Tensor,
-    random_int_tensor: torch.Tensor | None,
+    random_int_tensor: torch.Tensor,
     block_size: int = 2048,
     val_blocks: torch.Tensor | None = None,
 ) -> None:

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/util/scaled_optm.py

@@ -61,7 +61,7 @@ def adjust_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]:
     """
     # DoRA Scale (Magnitude Vector)
     if getattr(p, '_is_dora_scale', False):
-        return 0.0, cwd
+        return wd, cwd
 
     if getattr(p, '_is_oft', False):
         return wd, 0.0
@@ -76,7 +76,7 @@ def adjust_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]:
     else:
         # 1D Biases or generic 1D parameters
         # Centered WD safely regularizes the delta without collapsing base feature variance.
-        return 0.0, cwd
+        return wd, cwd
 
 
 def scale_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]:

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm/util/state_util.py

@@ -209,7 +209,7 @@ def fix_loaded_state_dtype(state: dict, p: torch.Tensor, group: dict) -> None:
 
     # Pre-define sets for known exact-match keys
     uint8_keys = {'sign', 'sign_slow', 'sign_buf', 'shifter'}
-    fp32_keys = {'mu_m_nmf', 'mv_m_nmf', 'mu_v_nmf', 'mv_v_nmf', 'mu_m_slow_nmf', 'mv_m_slow_nmf'}
+    fp32_keys = {'mu_m_nmf', 'mv_m_nmf', 'mu_v_nmf', 'mv_v_nmf', 'mu_m_slow_nmf', 'mv_m_slow_nmf', "mu_mbuf_nmf", "mv_mbuf_nmf", "mu_b_nmf", "normuon_v"}
 
     for key, val in state.items():
         if not isinstance(val, torch.Tensor):

{adv_optm-2.4.dev24 → adv_optm-2.5}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-2.4.dev24 → adv_optm-2.5}/setup.py

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

adv_optm-2.4.dev24/adv_optm/util/OrthoGrad.py

@@ -1,80 +0,0 @@
-import torch
-import math
-
-def _orthogonalize_gradient(p: torch.Tensor, grad: torch.Tensor) -> torch.Tensor:
-    """
-    Projects the gradient `grad` to be orthogonal to the parameter `p`.
-    Modified from:
-    https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability/blob/720d2444df12b851d6cb417ab08cf125c822b2ae/orthograd.py
-    """
-    original_shape = grad.shape
-    original_dtype = grad.dtype
-    w = p.view(-1).float()
-    g = grad.view(-1).float()
-    w_norm_sq = torch.dot(w, w).add_(1e-30)
-    proj = torch.dot(w, g) / w_norm_sq
-    g_orth = g.sub(w * proj)
-    g_norm = g.norm(2)
-    g_orth_norm = g_orth.norm(2).add_(1e-30)
-    g_orth_scaled = g_orth * (g_norm / g_orth_norm)
-    return g_orth_scaled.view(original_shape).to(original_dtype)
-
-
-def iterative_ortho_project(p: torch.Tensor, update: torch.Tensor, iters: int = 5) -> torch.Tensor:
-    """
-    Applies iterative alternating orthogonal projection to a 2D matrix.
-    Projects the update to be orthogonal to the parameter matrix along
-    rows and columns sequentially, alternating dimensions.
-    Inspired from Sinkhorn algorithm, 2 iterations is enough to converge
-    to cosine similarity of -1e4 to -1e-5 (semi orthogonal).
-    """
-    # 1D Vector Case fallback to the standard OrthoGrad
-    is_vector = p.ndim < 2 or getattr(p, '_is_dora_scale', False) or getattr(p, 'is_vector', False)
-    if is_vector:
-        return _orthogonalize_gradient(p, update)
-
-    original_shape = update.shape
-
-    # 2D+ Matrix Case
-    update_2d = update.view(update.shape[0], -1)
-    param_2d = p.view(p.shape[0], -1)
-
-    m, n = update_2d.shape
-
-    # Dynamically determine the order based on aspect ratio
-    row_first = m > n
-    dim = 0 if row_first else 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)
-
-    for _ in range(iters):
-        # First dimension
-        update_2d = _ortho_normed_dim(param_2d, update_2d, p_norm_sq_dim, dim)
-        # Second dimension
-        update_2d = _ortho_normed_dim(param_2d, update_2d, p_norm_sq_adim, 1 - dim)
-
-    return update_2d.view(original_shape)
-
-
-def _ortho_normed_dim(p_2d: torch.Tensor, update_2d: torch.Tensor, p_norm_sq: torch.Tensor, dim: int) -> torch.Tensor:
-    """
-    Projects the update to be orthogonal to p along 'dim' and dynamically restores 
-    the original magnitude of that dimension pre-projection.
-    """
-    # Record target magnitude before projection
-    norm_lb = 1 / math.sqrt(update_2d.shape[dim])
-    target_norm = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(norm_lb)
-
-    # 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)
-    # Standard gamma is -1, but -1.01 proved to converge faster
-    update_2d.addcmul_(proj, p_2d, value=-1.01)
-
-    # Magnitude Preservation
-    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)