adv-optm 2.4.dev2__tar.gz → 2.4.dev5__tar.gz

{adv_optm-2.4.dev2 → adv_optm-2.4.dev5}/PKG-INFO

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

{adv_optm-2.4.dev2 → adv_optm-2.4.dev5}/adv_optm/__init__.py

@@ -22,4 +22,4 @@ __all__ = [
     "SignSGD_adv",
 ]
 
-__version__ = "2.4.dev2"
+__version__ = "2.4.dev5"

{adv_optm-2.4.dev2 → adv_optm-2.4.dev5}/adv_optm/optim/AdaMuon_adv.py

@@ -206,6 +206,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
         if spectral_normalization and rms_rescaling:
             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.")
 
         defaults = {
             "lr": lr, "betas": betas, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
@@ -260,6 +262,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
             if group.get('use_muon') is None: # Fallback
                  group['use_muon'] = group.get('optim_type') == 'muon'
 
+        self.init_step()
+
         self.kourkoutas_helper = None
         if any(group.get('adam_kourkoutas_beta', False) for group in self.param_groups):
             self.kourkoutas_helper = KourkoutasHelper(self)
@@ -280,8 +284,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
     def load_state_dict(self, state_dict: dict) -> None:
         """
         Overrides default load_state_dict to implement a workaround for PyTorch's
-        automatic dtype casting. It ensures factorized states remain float32 for 
-        stability, preserves integer/float8 quantized anchor states, and forces 
+        automatic dtype casting. It ensures factorized states remain float32 for
+        stability, preserves integer/float8 quantized anchor states, and forces
         standard states onto the parameter's current dtype/device.
         """
         super().load_state_dict(state_dict)
@@ -419,7 +423,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             step_size = group['lr'] / bias_correction1
 
             if is_compiled:
-                step_size = torch.as_tensor(step_size, dtype=torch.float64)
+                step_size = torch.as_tensor(step_size)
                 adam_step_param = self._compiled_adam_step_parameter
             else:
                 adam_step_param = Muon_AuxAdam._adam_step_parameter
@@ -430,7 +434,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
         else: # Muon path
             if is_compiled:
-                lr = torch.as_tensor(group['lr'], dtype=torch.float64)
+                lr = torch.as_tensor(group['lr'])
                 muon_step_param = self._compiled_muon_step_parameter
             else:
                 lr = group['lr']
@@ -467,7 +471,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             else:
                 shape_for_scaling = p.shape
 
-            scaled_eps, adaptive_eps, spectral_target, wd_scale = get_spectral_scaling(shape_for_scaling, group['n_layers'])
+            scaled_eps, adaptive_eps, spectral_target, wd_scale = get_spectral_scaling(p, shape_for_scaling, group['n_layers'])
 
             weight_decay = group['weight_decay'] * wd_scale
             decoupled_wd = True

{adv_optm-2.4.dev2 → adv_optm-2.4.dev5}/adv_optm/optim/AdamW_adv.py

@@ -6,7 +6,7 @@ from typing import Optional, Callable
 
 from ..util import param_update
 from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state
-from ..util.update_util import _grams_update, _cautious_update
+from ..util.update_util import _grams_update, _cautious_update, _init_fisher_wd_scaler, _get_fisher_wd_scaler
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
@@ -29,6 +29,9 @@ class AdamW_adv(torch.optim.Optimizer):
         eps (float): term added to the denominator to improve
             numerical stability (default: 1e-8)
         weight_decay (float): weight decay (L2 penalty) (default: 0).
+        fisher_wd (bool): whether to use Fisher Adam (FAdam) weight decay, mapping
+            the decay direction through the empirical Fisher information matrix and
+            clipping its RMS. (default: False)
         cautious_wd (bool): Enables Cautious Weight Decay. If True, weight decay is
             applied only to parameter coordinates where the sign of the parameter
             and the sign of the optimizer update align (default: False).
@@ -91,7 +94,7 @@ class AdamW_adv(torch.optim.Optimizer):
             'int4': Uses 4-bit block-wise quantization (block size 32).
         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) 
+        factored_2nd (bool): whether to keep the first moment uncompressed (dense)
             while only factorizing the second moment. (default: True)
     """
 
@@ -103,6 +106,7 @@ class AdamW_adv(torch.optim.Optimizer):
         eps: float = 1e-8,
         # Decoupled/cautious weight decay
         weight_decay: float = 0.0,
+        fisher_wd: bool = False,
         cautious_wd: bool = False,
         # Adam's Bias Correction
         use_bias_correction: bool = True,
@@ -155,7 +159,8 @@ class AdamW_adv(torch.optim.Optimizer):
             cautious_mask = False
 
         defaults = {
-            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
+            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
             "use_atan2": use_atan2,
             "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction,
             "beta3_ema": beta3_ema, "alpha": alpha, "compiled_optimizer": compiled_optimizer,
@@ -192,8 +197,8 @@ class AdamW_adv(torch.optim.Optimizer):
     def load_state_dict(self, state_dict: dict) -> None:
         """
         Overrides default load_state_dict to implement a workaround for PyTorch's
-        automatic dtype casting. It ensures factorized states remain float32 for 
-        stability, preserves integer/float8 quantized anchor states, and forces 
+        automatic dtype casting. It ensures factorized states remain float32 for
+        stability, preserves integer/float8 quantized anchor states, and forces
         standard states onto the parameter's current dtype/device.
         """
         super().load_state_dict(state_dict)
@@ -273,6 +278,8 @@ class AdamW_adv(torch.optim.Optimizer):
 
             _init_anchor(p, state, group)
 
+            _init_fisher_wd_scaler(group, state, p)
+
         beta1, beta2 = group['betas']
 
         current_step = state['step']
@@ -294,7 +301,7 @@ class AdamW_adv(torch.optim.Optimizer):
         random_int_tensor = None
 
         if group.get('compiled_optimizer', False):
-            step_size = torch.as_tensor(step_size, dtype=torch.float64)
+            step_size = torch.as_tensor(step_size)
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
@@ -349,7 +356,11 @@ class AdamW_adv(torch.optim.Optimizer):
                     update_mt = mt if not factored_2nd else mt.clone()
 
             vt = _reconstruct_state((state['mu_v_nmf'], state['mv_v_nmf']), signed=False)
-            vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
+
+            if isinstance(beta2, torch.Tensor) and beta2.dim() > 0:
+                vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped * (1.0 - beta2))
+            else:
+                vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
 
             if self.use_AdEMAMix:
                 if factored_2nd:
@@ -363,7 +374,7 @@ class AdamW_adv(torch.optim.Optimizer):
                     update = update_mt.add_(mt_slow, alpha=alpha)
                 else:
                     update = grad_reshaped.add(mt_slow, alpha=alpha)
-                
+
                 if not factored_2nd:
                     # Factorize
                     state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'] = _factorize_state(mt_slow, signed=True)
@@ -385,6 +396,9 @@ class AdamW_adv(torch.optim.Optimizer):
                 denom = vt.sqrt_()
                 denom.div_(sqrt_bias_correction2).add_(group['eps'])
                 update.div_(denom)
+
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, group['use_atan2'])
+
             del vt
 
             update = update.view(p.shape)
@@ -413,7 +427,10 @@ class AdamW_adv(torch.optim.Optimizer):
                 update = update_mt if beta1 > 0 else grad.clone()
 
             exp_avg_sq = state['exp_avg_sq']
-            exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+            if isinstance(beta2, torch.Tensor) and beta2.dim() > 0:
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad * (1.0 - beta2))
+            else:
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
 
             if group['use_atan2']:
                 denom = exp_avg_sq.sqrt()
@@ -423,6 +440,9 @@ class AdamW_adv(torch.optim.Optimizer):
                 denom = exp_avg_sq.sqrt()
                 denom.div_(sqrt_bias_correction2).add_(group['eps'])
                 update.div_(denom)
+
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, group['use_atan2'])
+
             del denom
 
         update_scaling = step_size * A if group['use_atan2'] else step_size
@@ -431,7 +451,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)
+        param_update.apply_parameter_update(self, p, group, update, step_size, 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.dev2 → adv_optm-2.4.dev5}/adv_optm/optim/Adopt_adv.py

@@ -7,7 +7,7 @@ from ..util import param_update
 from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state, _nnmf
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
-from ..util.update_util import _grams_update, _cautious_update, _scale_sim_AdEMAMix_update
+from ..util.update_util import _grams_update, _cautious_update, _scale_sim_AdEMAMix_update, _init_fisher_wd_scaler, _get_fisher_wd_scaler
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
 from ..util.centered_decay import _init_anchor
 
@@ -33,6 +33,9 @@ class Adopt_adv(torch.optim.Optimizer):
         eps (float): term added to the denominator to improve
             numerical stability (default: 1e-6)
         weight_decay (float): weight decay (L2 penalty) (default: 0)
+        fisher_wd (bool): whether to use Fisher Adam (FAdam) weight decay, mapping
+            the decay direction through the empirical Fisher information matrix and
+            clipping its RMS. (default: False)
         cautious_wd (bool): Enables Cautious Weight Decay. If True, weight decay is
             applied only to parameter coordinates where the sign of the parameter
             and the sign of the optimizer update align (default: False).
@@ -107,7 +110,7 @@ class Adopt_adv(torch.optim.Optimizer):
             'int4': Uses 4-bit block-wise quantization (block size 32).
         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) 
+        factored_2nd (bool): whether to keep the first moment uncompressed (dense)
             while only factorizing the second moment. (default: True)
     """
 
@@ -119,6 +122,7 @@ class Adopt_adv(torch.optim.Optimizer):
         eps: float = 1e-6,
         # Decoupled/cautious weight decay
         weight_decay: float = 0.0,
+        fisher_wd: bool = False,
         cautious_wd: bool = False,
         # ADOPT clipping
         clip_lambda: Optional[Callable[[int], float]] = lambda step: step**0.25,
@@ -181,7 +185,8 @@ class Adopt_adv(torch.optim.Optimizer):
             print("Warning: cautious is incompatible with Simplified_AdEMAMix, Disabling cautious.")
 
         defaults = {
-            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
+            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
             "beta3_ema": beta3_ema, "alpha": alpha,
             "alpha_grad": alpha_grad,
             "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
@@ -189,7 +194,7 @@ class Adopt_adv(torch.optim.Optimizer):
             "scaled_optm": scaled_optm,
             "centered_wd": centered_wd,
             "centered_wd_mode": centered_wd_mode,
-            "nnmf_factor": nnmf_factor, "vector_reshape": vector_reshape, "factored_2nd": factored_2nd, 
+            "nnmf_factor": nnmf_factor, "vector_reshape": vector_reshape, "factored_2nd": factored_2nd,
             "compiled_optimizer": compiled_optimizer,
         }
         self.clip_lambda = clip_lambda
@@ -222,8 +227,8 @@ class Adopt_adv(torch.optim.Optimizer):
     def load_state_dict(self, state_dict: dict) -> None:
         """
         Overrides default load_state_dict to implement a workaround for PyTorch's
-        automatic dtype casting. It ensures factorized states remain float32 for 
-        stability, preserves integer/float8 quantized anchor states, and forces 
+        automatic dtype casting. It ensures factorized states remain float32 for
+        stability, preserves integer/float8 quantized anchor states, and forces
         standard states onto the parameter's current dtype/device.
         """
         super().load_state_dict(state_dict)
@@ -244,6 +249,19 @@ class Adopt_adv(torch.optim.Optimizer):
         grad = p.grad
         state = self.state[p]
 
+
+        beta1, beta2 = group['betas']
+
+        if group.get('kourkoutas_beta', False):
+            if 'step' not in state:
+                current_step = 0
+            else:
+                current_step = state['step']
+            # Call prepare_step() once at the beginning of the step for all params
+            self.kourkoutas_helper.maybe_prepare_step(current_step, p.device)
+            # Get the dynamic beta2 calculated in prepare_step()
+            beta2 = self.kourkoutas_helper.get_beta2(p, group)
+
         # State Initialization
         if 'step' not in state:
             state['step'] = 0
@@ -256,6 +274,12 @@ class Adopt_adv(torch.optim.Optimizer):
 
             dtype = torch.float32 if state['factored'] else p.dtype
 
+            vt_init = grad.pow(2).to(dtype)
+            if isinstance(beta2, torch.Tensor) and beta2.dim() > 0:
+                vt_init.mul_(beta2).addcmul_(grad.to(dtype), grad.to(dtype) * (1.0 - beta2))
+            else:
+                vt_init.mul_(beta2).addcmul_(grad.to(dtype), grad.to(dtype), value=1.0 - beta2)
+
             if state['factored']:
                 state['effective_shape'] = _get_effective_shape(p.numel())
                 d1, d2 = state['effective_shape']
@@ -279,33 +303,23 @@ class Adopt_adv(torch.optim.Optimizer):
                     if self.use_AdEMAMix:
                         state['exp_avg_slow'] = torch.zeros_like(p, device=p.device, dtype=dtype)
                 # Second moment (v)
-                vt_init = grad.to(dtype).view(d1, d2).square()
-                # Allocate NMF factors for vt
-                state['mu_v_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype)
-                state['mv_v_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
-                # Initialize v_0
-                state['mu_v_nmf'], state['mv_v_nmf'] = _nnmf(vt_init)
+                state['mu_v_nmf'], state['mv_v_nmf'] = _nnmf(vt_init.view(d1, d2))
                 del vt_init
             else: # Fallback for non-factored tensors
                 if group['betas'][0] > 0:
                     state['exp_avg'] = torch.zeros_like(p, device=p.device, dtype=dtype)
                 if self.use_AdEMAMix:
                     state['exp_avg_slow'] = torch.zeros_like(p, device=p.device, dtype=dtype)
-                state['exp_avg_sq'] = grad.to(dtype).square()
+                state['exp_avg_sq'] = vt_init
 
             if group.get('scaled_optm', False) and is_spectral(p):
                 init_spectral_norm(group, state, p)
 
             _init_anchor(p, state, group)
 
-        beta1, beta2 = group['betas']
+            _init_fisher_wd_scaler(group, state, p)
 
         current_step = state['step']
-        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, p.device)
-            # Get the dynamic beta2 calculated in prepare_step()
-            beta2 = self.kourkoutas_helper.get_beta2(p, group)
 
         # 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:
@@ -315,7 +329,7 @@ class Adopt_adv(torch.optim.Optimizer):
         random_int_tensor = None
 
         if group.get('compiled_optimizer', False):
-            lr = torch.as_tensor(group['lr'], dtype=torch.float64)
+            lr = torch.as_tensor(group['lr'])
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
@@ -359,9 +373,13 @@ class Adopt_adv(torch.optim.Optimizer):
 
             # ADOPT Step A: Decorrelate g_t using v_{t-1}
             denom = vt.sqrt()
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, self.use_atan2)
 
             # Update second moment v_t for the *next* step using raw g_t
-            vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
+            if isinstance(beta2, torch.Tensor) and beta2.dim() > 0:
+                vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped * (1.0 - beta2))
+            else:
+                vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
             # Factorize
             state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(vt, signed=False)
             del vt
@@ -434,6 +452,7 @@ class Adopt_adv(torch.optim.Optimizer):
 
             # ADOPT Step A: Decorrelate g_t using v_{t-1}
             denom = vt.sqrt()
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, self.use_atan2)
 
             if self.use_atan2:
                 normalized_grad = torch.atan2(grad, denom, out=denom)
@@ -475,9 +494,11 @@ class Adopt_adv(torch.optim.Optimizer):
                 else:
                     update = normalized_grad
 
-
             # Update second moment v_t for the next step using raw g_t
-            vt.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+            if isinstance(beta2, torch.Tensor) and beta2.dim() > 0:
+                vt.mul_(beta2).addcmul_(grad, grad * (1.0 - beta2))
+            else:
+                vt.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
 
         update_scaling = lr * A if self.use_atan2 else lr
 
@@ -487,7 +508,7 @@ class Adopt_adv(torch.optim.Optimizer):
             update.mul_(update_scaling)
 
         # Parameter Update
-        param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
+        param_update.apply_parameter_update(self, p, group, update, 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.dev2 → adv_optm-2.4.dev5}/adv_optm/optim/Lion_Prodigy_adv.py

@@ -225,8 +225,8 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
-            d = torch.as_tensor(group['d'], dtype=torch.float64)
-            dlr = torch.as_tensor(dlr, dtype=torch.float64)
+            d = torch.as_tensor(group['d'])
+            dlr = torch.as_tensor(dlr)
             step_param_fn = self._compiled_step_parameter
         else:
             d = group['d']

{adv_optm-2.4.dev2 → adv_optm-2.4.dev5}/adv_optm/optim/Lion_adv.py

@@ -8,6 +8,8 @@ from ..util.factorization_util import _get_effective_shape, _reconstruct_state,
 from ..util.lion_k import _get_lion_k_update
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
 from ..util.centered_decay import _init_anchor
+from ..util.update_util import _get_l1_adaptive_lr
+from ..util.signed_util import apply_stochastic_sign
 
 
 class Lion_adv(torch.optim.Optimizer):
@@ -44,9 +46,10 @@ class Lion_adv(torch.optim.Optimizer):
             parameter dimensionality. Sets p=2.0 for 4D tensors (Conv2D) (Biases/Norms) to
             use Spherical updates, and p=1.0 for others (Linear/Embeddings) to use Sign
             updates. Overrides explicit kappa_p value. (default: False).
+        stochastic_sign (bool): whether to use the Stochastic Sign operator. (default: False)
         freeze_on_flip (bool): Projected SignGD One-hit freeze. Masks updates for
             coordinates where the gradient sign flips compared to the previous step. (default: False)
-        l1_adaptive (bool): Scales learning rate dynamically 
+        l1_adaptive (bool): Scales learning rate dynamically
             by the L1 norm of the gradient to handle gradient heterogeneity. (default: False).
         centered_wd (float): Centered Weight Decay coefficient. Instead of decaying weights
             toward zero, they are decayed toward their initial values (anchors). This
@@ -79,6 +82,8 @@ class Lion_adv(torch.optim.Optimizer):
         # Lion-k
         kappa_p: float = 1.0,
         auto_kappa_p: bool = False,
+        # Stochastic Sign Operator
+        stochastic_sign: bool = False,
         # Projected and adaptive sign
         freeze_on_flip: bool = False,
         l1_adaptive: bool = False,
@@ -110,6 +115,7 @@ class Lion_adv(torch.optim.Optimizer):
             clip_threshold=clip_threshold,
             kappa_p=kappa_p,
             auto_kappa_p=auto_kappa_p,
+            stochastic_sign=stochastic_sign,
             freeze_on_flip=freeze_on_flip,
             l1_adaptive=l1_adaptive,
             scaled_optm= scaled_optm,
@@ -137,8 +143,8 @@ class Lion_adv(torch.optim.Optimizer):
     def load_state_dict(self, state_dict: dict) -> None:
         """
         Overrides default load_state_dict to implement a workaround for PyTorch's
-        automatic dtype casting. It ensures factorized states remain float32 for 
-        stability, preserves integer/float8 quantized anchor states, and forces 
+        automatic dtype casting. It ensures factorized states remain float32 for
+        stability, preserves integer/float8 quantized anchor states, and forces
         standard states onto the parameter's current dtype/device.
         """
         super().load_state_dict(state_dict)
@@ -201,19 +207,22 @@ class Lion_adv(torch.optim.Optimizer):
         lr = group["lr"]
 
         random_int_tensor = None
+        random_noise_tensor = None
 
         if group.get('compiled_optimizer', False):
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
-            lr = torch.as_tensor(lr, dtype=torch.float64)
+            if group.get('stochastic_sign', False):
+                random_noise_tensor = param_update._get_random_noise_for_sso(p)
+            lr = torch.as_tensor(lr)
             step_param_fn = self._compiled_step_parameter
         else:
             step_param_fn = self._step_parameter
 
-        step_param_fn(p, grad, state, group, lr, random_int_tensor)
+        step_param_fn(p, grad, state, group, lr, random_int_tensor, random_noise_tensor)
 
-    def _step_parameter(self, p, grad, state, group, lr, random_int_tensor):
+    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["clip_threshold"] > 0.0:
@@ -251,9 +260,6 @@ class Lion_adv(torch.optim.Optimizer):
             # Compute update term c_t
             update = torch.lerp(grad_reshaped, exp_avg, beta1)
 
-            if group.get("l1_adaptive", False) and kappa_p == 1:
-                lr = lr * (update.norm(p=1))
-
             # Standard Lion momentum update
             # m_t = beta2 * m_{t-1} + (1-beta2) * g_t
             exp_avg.lerp_(grad_reshaped, 1 - beta2)
@@ -262,7 +268,12 @@ class Lion_adv(torch.optim.Optimizer):
             state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(exp_avg, signed=True)
             del exp_avg
 
-            update = _get_lion_k_update(update, kappa_p)
+            if freeze_on_flip:
+                # Fast binary diff (XOR) from momentum sign directly
+                flipped_packed = prev_sign_packed ^ state['sign']
+                flipped_mask = _unpack_bools(flipped_packed, original_m=d2).view_as(update)
+                update = torch.where(flipped_mask, 0.0, update)
+                del prev_sign_packed, flipped_packed, flipped_mask
 
             if self.cautious_mask:
                 mask = (update * grad_reshaped > 0).to(grad_reshaped.dtype)
@@ -272,12 +283,12 @@ class Lion_adv(torch.optim.Optimizer):
 
             update = update.view(p.shape)
 
-            if freeze_on_flip:
-                # Fast binary diff (XOR) from momentum sign directly
-                flipped_packed = prev_sign_packed ^ state['sign']
-                flipped_mask = _unpack_bools(flipped_packed, original_m=d2).view_as(update)
-                update = torch.where(flipped_mask, 0.0, update)
-                del prev_sign_packed, flipped_packed, flipped_mask
+            l1_mean = _get_l1_adaptive_lr(p, update, state, group, kappa_p, rescale=False)
+
+            if group.get('stochastic_sign', False):
+                update = apply_stochastic_sign(update, noise=random_noise_tensor)
+            else:
+                update = _get_lion_k_update(update, kappa_p)
 
         else:
             # Fallback to standard Lion logic
@@ -286,10 +297,13 @@ class Lion_adv(torch.optim.Optimizer):
             # Compute update term
             update = torch.lerp(grad, exp_avg, beta1)
 
-            if group.get("l1_adaptive", False) and kappa_p == 1:
-                lr = lr * (update.norm(p=1))
+            # Standard Lion momentum update
+            exp_avg.lerp_(grad, 1 - beta2)
 
-            update = _get_lion_k_update(update, kappa_p)
+            if freeze_on_flip:
+                current_sign = (update > 0).to(torch.uint8)
+                update = torch.where(current_sign == state['prev_sign'], update, 0.0)
+                state['prev_sign'] = current_sign
 
             if self.cautious_mask:
                 mask = (update * grad > 0).to(grad.dtype)
@@ -297,20 +311,22 @@ class Lion_adv(torch.optim.Optimizer):
                 update.mul_(mask)
                 del mask
 
-            # Standard Lion momentum update
-            exp_avg.lerp_(grad, 1 - beta2)
+            l1_mean = _get_l1_adaptive_lr(p, update, state, group, kappa_p, rescale=False)
 
-            if freeze_on_flip:
-                current_sign = (update > 0).to(torch.uint8)
-                update = torch.where(current_sign == state['prev_sign'], update, 0.0)
-                state['prev_sign'] = current_sign
+            if group.get('stochastic_sign', False):
+                update = apply_stochastic_sign(update, noise=random_noise_tensor)
+            else:
+                update = _get_lion_k_update(update, kappa_p)
+
+        if l1_mean is not None:
+            update.mul_(l1_mean)
 
         if group.get('scaled_optm', False):
             update = scale_update(p, update, lr, vector_state=state.get('spectral_v'))
         else:
             update.mul_(lr)
 
-        param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
+        param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor, wd_scaler=l1_mean)
 
     def compile(self, *args, **kwargs):
         self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)

{adv_optm-2.4.dev2 → adv_optm-2.4.dev5}/adv_optm/optim/Muon_adv.py

@@ -183,6 +183,8 @@ class Muon_adv(torch.optim.Optimizer):
         if spectral_normalization and rms_rescaling:
             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.")
 
         defaults = {
             "lr": lr, "beta1": beta1, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
@@ -239,6 +241,8 @@ class Muon_adv(torch.optim.Optimizer):
             if group.get('use_muon') is None: # Fallback
                  group['use_muon'] = group.get('optim_type') == 'muon'
 
+        self.init_step()
+
         self.kourkoutas_helper = None
         if any(group.get('adam_kourkoutas_beta', False) for group in self.param_groups):
             self.kourkoutas_helper = KourkoutasHelper(self)
@@ -259,8 +263,8 @@ class Muon_adv(torch.optim.Optimizer):
     def load_state_dict(self, state_dict: dict) -> None:
         """
         Overrides default load_state_dict to implement a workaround for PyTorch's
-        automatic dtype casting. It ensures factorized states remain float32 for 
-        stability, preserves integer/float8 quantized anchor states, and forces 
+        automatic dtype casting. It ensures factorized states remain float32 for
+        stability, preserves integer/float8 quantized anchor states, and forces
         standard states onto the parameter's current dtype/device.
         """
         super().load_state_dict(state_dict)
@@ -393,7 +397,7 @@ class Muon_adv(torch.optim.Optimizer):
             step_size = group['lr'] / bias_correction1
 
             if is_compiled:
-                step_size = torch.as_tensor(step_size, dtype=torch.float64)
+                step_size = torch.as_tensor(step_size)
                 adam_step_param = self._compiled_adam_step_parameter
             else:
                 adam_step_param = Muon_AuxAdam._adam_step_parameter
@@ -404,7 +408,7 @@ class Muon_adv(torch.optim.Optimizer):
 
         else: # Muon path
             if is_compiled:
-                lr = torch.as_tensor(group['lr'], dtype=torch.float64)
+                lr = torch.as_tensor(group['lr'])
                 muon_step_param = self._compiled_muon_step_parameter
             else:
                 lr = group['lr']