adv-optm 2.4.dev12__tar.gz → 2.4.dev14__tar.gz

{adv_optm-2.4.dev12 → adv_optm-2.4.dev14}/PKG-INFO

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

{adv_optm-2.4.dev12 → adv_optm-2.4.dev14}/adv_optm/__init__.py

@@ -24,4 +24,4 @@ __all__ = [
     "SinkSGD_adv",
 ]
 
-__version__ = "2.4.dev12"
+__version__ = "2.4.dev14"

{adv_optm-2.4.dev12 → adv_optm-2.4.dev14}/adv_optm/optim/AdaMuon_adv.py

@@ -59,14 +59,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
         orthogonal_gradient (bool): whether to use OrthoGrad.  (default: False)
         nesterov (bool): enables Nesterov momentum (default: False).
         use_atan2 (bool): whether to use the atan2 update rule. (default: False)
-        Simplified_AdEMAMix (bool): whether to use the Simplified AdEMAMix update rule.
-            This changes the update  to `alpha_grad * grad + mt`, which can be
-            more responsive, especially for small batch sizes. (default: False)
-        alpha_grad (float): Mixing coefficient for the Simplified AdEMAMix update rule
-            (only used when `Simplified_AdEMAMix` is `True`). Controls the weight of the
-            current gradient. For small batch sizes, use high values (e.g., 10-100) to be
-            more responsive. For large batch sizes, use low values (e.g., 0-1) for
-            stability. (default: 100.0)
         vector_reshape (bool): whether to reshape 1D vectors into 2D
             matrices to apply low-rank compression (default: True).
         kappa_p (float, optional): The p-value for the update geometry  (domain [1.0, 2.0]).
@@ -117,6 +109,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         adam_betas (tuple[float, float]): Betas for the AdamW optimizer part.
         adam_eps (float): Epsilon for the AdamW optimizer part.
         adam_weight_decay (float): Weight decay for the AdamW optimizer part.
+        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_cautious_mask (bool): Cautious masking for AdamW.
@@ -125,8 +118,17 @@ class AdaMuon_adv(torch.optim.Optimizer):
         adam_use_AdEMAMix (bool): AdEMAMix for AdamW.
         adam_beta3_ema (float): Beta3 for AdEMAMix.
         adam_alpha (float): Alpha for AdEMAMix.
+        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.
+        adam_beta2_min (float): Minimum beta2 for Kourkoutas-β. (default: 0.9)
+        adam_ema_alpha (float): EMA alpha for Kourkoutas-β. (default: 0.95)
+        adam_tiny_spike (float): Tiny spike for Kourkoutas-β. (default: 1e-9)
+        adam_k_warmup_steps (int): Warmup steps for Kourkoutas-β. (default: 0)
+        adam_spectral_normalization (bool): Enable explicit spectral normalization for AdamW. (default: False)
+        adam_state_precision (str): Precision for AuxAdam states. Options: 'auto', 'fp32', 'bf16_sr', 'fp8_sr', 'int8_sr', 'factored'. (default: 'auto')
         adam_nnmf_factor (bool): 1-bit factored for AdamW.
+        adam_factored_2nd (bool): Factorize only the second moment (v_t) for AuxAdam. (default: False)
     """
 
     def __init__(
@@ -140,6 +142,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         cautious_wd: bool = False,
         # Nesterov momentum
         nesterov: bool = True,
+        nesterov_coef: float | None = None,
         # RMS Rescaling
         rms_rescaling: bool = True,
         # Newton Schulz
@@ -152,9 +155,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
         orthogonal_gradient: bool = False,
         # Adam_atan2 (scale invariant)
         use_atan2: bool = False,
-        # One-EMA AdEMAMix
-        Simplified_AdEMAMix: bool = False,
-        alpha_grad: float = 100.0,
         # NorMuon
         normuon_variant: bool = False,
         # Boolean to spilt param
@@ -190,6 +190,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         adam_betas: tuple[float, float] = (0.9, 0.99),
         adam_eps: float | None = 1e-8,
         adam_weight_decay: float = 0.0,
+        adam_fisher_wd: bool = False,
         adam_use_bias_correction: bool = True,
         adam_use_atan2: bool = False,
         adam_cautious_mask: bool = False,
@@ -198,12 +199,17 @@ class AdaMuon_adv(torch.optim.Optimizer):
         adam_use_AdEMAMix: bool = False,
         adam_beta3_ema: float = 0.9999,
         adam_alpha: float = 5.0,
+        adam_nesterov: bool = False,
+        adam_nesterov_coef: float | None = None,
         adam_kourkoutas_beta: bool = False,
         adam_beta2_min: float = 0.9,
         adam_ema_alpha: float = 0.95,
         adam_tiny_spike: float = 1e-9,
         adam_k_warmup_steps: int = 0,
+        adam_spectral_normalization: bool = False,
+        adam_state_precision: str = "auto",
         adam_nnmf_factor: bool = False,
+        adam_factored_2nd: bool = False,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -211,9 +217,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
         if not (ns_steps > 0):
             raise ValueError(f"Newton-Schulz steps should be > 0. Got {ns_steps}")
-        if Simplified_AdEMAMix and nesterov:
-            print("Warning: nesterov is incompatible with Simplified_AdEMAMix, Disabling nesterov.")
-            nesterov = False
         if spectral_normalization and rms_rescaling:
             print("Warning: spectral_normalization is incompatible with rms_rescaling, Disabling rms_rescaling.")
             rms_rescaling = False
@@ -221,17 +224,20 @@ class AdaMuon_adv(torch.optim.Optimizer):
             ValueError("spectral_normalization violates accelerated Newton-Schulz assumptions. Pick one of them.")
 
         state_precision = state_precision.lower()
-        valid_precisions = {"auto", "fp32", "bf16_sr", "fp8_sr", "int8_sr"}
+        valid_precisions = {"auto", "fp32", "factored", "bf16_sr", "fp8_sr", "int8_sr"}
         if state_precision not in valid_precisions:
             raise ValueError(f"state_precision must be one of {valid_precisions}. Got {state_precision}")
 
+        adam_state_precision = adam_state_precision.lower()
+        if adam_state_precision not in valid_precisions:
+            raise ValueError(f"adam_state_precision must be one of {valid_precisions}. Got {adam_state_precision}")
+
         defaults = {
             "lr": lr, "betas": betas, "weight_decay": weight_decay, "cautious_wd": cautious_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, "use_atan2":use_atan2,
-            "Simplified_AdEMAMix": Simplified_AdEMAMix, "alpha_grad": alpha_grad,
+            "nesterov":nesterov, "nesterov_coef": nesterov_coef, "use_atan2":use_atan2,
             "normuon_variant": normuon_variant, "orthogonal_gradient": orthogonal_gradient,
             "compiled_optimizer":compiled_optimizer,
             "use_muon": use_muon,
@@ -254,13 +260,18 @@ class AdaMuon_adv(torch.optim.Optimizer):
             "centered_wd_mode": centered_wd_mode,
             # AdamW_adv defaults
             "adam_betas": adam_betas, "adam_eps": adam_eps, "adam_weight_decay": adam_weight_decay,
+            "adam_fisher_wd": adam_fisher_wd,
             "adam_use_bias_correction": adam_use_bias_correction, "adam_use_atan2": adam_use_atan2,
             "adam_cautious_mask": adam_cautious_mask, "adam_grams_moment": adam_grams_moment,
             "adam_orthogonal_gradient": adam_orthogonal_gradient,
             "adam_use_AdEMAMix": adam_use_AdEMAMix, "adam_beta3_ema": adam_beta3_ema, "adam_alpha": adam_alpha,
+            "adam_nesterov": adam_nesterov, "adam_nesterov_coef": adam_nesterov_coef,
             "adam_kourkoutas_beta": adam_kourkoutas_beta, "adam_beta2_min": adam_beta2_min,
             "adam_ema_alpha": adam_ema_alpha, "adam_tiny_spike": adam_tiny_spike,
-            "adam_k_warmup_steps": adam_k_warmup_steps, "adam_nnmf_factor": adam_nnmf_factor,
+            "adam_k_warmup_steps": adam_k_warmup_steps,
+            "adam_spectral_normalization": adam_spectral_normalization,
+            "adam_state_precision": adam_state_precision,
+            "adam_nnmf_factor": adam_nnmf_factor, "adam_factored_2nd": adam_factored_2nd,
         }
         self.stochastic_rounding = stochastic_rounding
         self._init_lr = lr
@@ -447,13 +458,24 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
             step_size = group['lr'] / bias_correction1
 
+            random_int_state_tensor = None
             if is_compiled:
                 step_size = torch.as_tensor(step_size)
                 adam_step_param = self._compiled_adam_step_parameter
+                
+                # Generate state SR random tensor when compiled
+                actual_precision = group.get('adam_actual_state_precision', 'auto')
+                random_int_state_tensor = random_int_tensor
+                if actual_precision == 'bf16_sr' and random_int_state_tensor is None:
+                    random_int_state_tensor = param_update._get_random_int_for_sr(p)
+                elif actual_precision == 'int8_sr':
+                    random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
+                elif actual_precision == 'fp8_sr':
+                    random_int_state_tensor = param_update._get_random_int_for_fp8_sr(p)
             else:
                 adam_step_param = Muon_AuxAdam._adam_step_parameter
 
-            adam_step_param(self, p, grad, state, group, beta1_adam, beta2_adam, sqrt_bias_correction2, step_size, random_int_tensor)
+            adam_step_param(self, p, grad, state, group, beta1_adam, beta2_adam, sqrt_bias_correction2, step_size, random_int_tensor, random_int_state_tensor)
 
             state['step'] += 1
 
@@ -465,7 +487,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 # Generate state SR random tensor when compiled
                 actual_precision = group['actual_state_precision']
                 random_int_state_tensor = random_int_tensor
-                if actual_precision == 'bf16_sr' and random_int_state_tensor is not None:
+                if actual_precision == 'bf16_sr' and random_int_state_tensor is None:
                     random_int_state_tensor = param_update._get_random_int_for_sr(p)
                 elif actual_precision == 'int8_sr':
                     random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
@@ -488,8 +510,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         grad = upcast_grad_for_precision(grad, state, group.get('state_precision', 'auto'))
         beta1, beta2 = group['betas']
         nesterov = group['nesterov']
-        Simplified_AdEMAMix = group['Simplified_AdEMAMix']
-        alpha_grad = group['alpha_grad']
+        nesterov_coef = group.get('nesterov_coef', None)
 
         # Update geometry
         kappa_p = group.get("kappa_p", 1.0)
@@ -513,7 +534,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
         # MARS-M Approximated (Variance Reduction)
         if group.get('approx_mars', False):
-            grad = approx_mars(grad, state['last_grad'], group['mars_gamma'], beta1, Simplified_AdEMAMix=Simplified_AdEMAMix)
+            grad = approx_mars(grad, state['last_grad'], group['mars_gamma'], beta1)
 
 
         if group.get("orthogonal_gradient"):
@@ -527,15 +548,11 @@ class AdaMuon_adv(torch.optim.Optimizer):
             mt_buf = _reconstruct_state((state['mu_mbuf_nmf'], state['mv_mbuf_nmf'], state['sign_buf'], d2), signed=True)
 
             # Update momentum in full-size
-            if not Simplified_AdEMAMix:
-                mt_buf.lerp_(grad_reshaped, 1 - beta1)
-            else:
-                mt_buf.mul_(beta1).add_(grad_reshaped)
+            mt_buf.lerp_(grad_reshaped, 1 - beta1)
 
             if nesterov:
-                update = grad_reshaped.lerp(mt_buf, beta1)
-            elif Simplified_AdEMAMix:
-                update = torch.add(mt_buf, grad_reshaped, alpha=alpha_grad)
+                nv_coef = beta1 if nesterov_coef is None else nesterov_coef
+                update = grad_reshaped.lerp(mt_buf, nv_coef)
             else:
                 update = mt_buf.clone()
 
@@ -586,15 +603,11 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
             # Momentum update
             mt_buf = get_state(state, 'momentum_buffer', actual_precision)
-            if not Simplified_AdEMAMix:
-                mt_buf.lerp_(grad, 1 - beta1)
-            else:
-                mt_buf.mul_(beta1).add_(grad)
+            mt_buf.lerp_(grad, 1 - beta1)
 
             if nesterov:
-                update = grad.lerp(mt_buf, beta1)
-            elif Simplified_AdEMAMix:
-                update = mt_buf.add(grad, alpha=alpha_grad)
+                nv_coef = beta1 if nesterov_coef is None else nesterov_coef
+                update = grad.lerp(mt_buf, nv_coef)
             else:
                 update = mt_buf.clone()
 

{adv_optm-2.4.dev12 → adv_optm-2.4.dev14}/adv_optm/optim/AdamW_adv.py

@@ -28,7 +28,8 @@ class AdamW_adv(torch.optim.Optimizer):
         betas (tuple[float, float]): coefficients used for computing running
             averages of gradient and its square (default: (0.9, 0.999))
         eps (float): term added to the denominator to improve
-            numerical stability (default: 1e-8)
+            numerical stability. Set to None for scale invariant eps (vector
+            lower bound) (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
@@ -127,6 +128,9 @@ class AdamW_adv(torch.optim.Optimizer):
         use_AdEMAMix: bool = False,
         beta3_ema: float = 0.9999,
         alpha: float = 5.0,
+        # Nesterov momentum
+        nesterov: bool = False,
+        nesterov_coef: float | None = None,
         # K-b (adaptive beta2)
         kourkoutas_beta: bool = False,
         beta2_min: float = 0.9,
@@ -176,7 +180,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,
-            "use_atan2": use_atan2,
+            "use_atan2": use_atan2, "nesterov": nesterov, "nesterov_coef": nesterov_coef,
             "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction,
             "beta3_ema": beta3_ema, "alpha": alpha, "compiled_optimizer": compiled_optimizer,
             "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
@@ -195,6 +199,8 @@ class AdamW_adv(torch.optim.Optimizer):
         self._init_lr = lr
         super().__init__(params, defaults)
 
+        self.init_step()
+
         if self.kourkoutas_beta:
             self.kourkoutas_helper = KourkoutasHelper(self)
 
@@ -232,12 +238,13 @@ class AdamW_adv(torch.optim.Optimizer):
     def supports_flat_params(self):
         return False
 
-    @torch.no_grad()
-    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
-        if p.grad is None:
-            return
+    def init_step(self):
+        for group in self.param_groups:
+            for i, p in enumerate(group['params']):
+                self.__init_state(p, group)
 
-        grad = p.grad
+    @torch.no_grad()
+    def __init_state(self, p, group):
         state = self.state[p]
 
         # State Initialization
@@ -303,6 +310,15 @@ class AdamW_adv(torch.optim.Optimizer):
 
             _init_fisher_wd_scaler(group, state, p)
 
+    @torch.no_grad()
+    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
+        if p.grad is None:
+            return
+
+        grad = p.grad
+        state = self.state[p]
+        self.__init_state(p, group)
+
         beta1, beta2 = group['betas']
 
         current_step = state['step']
@@ -353,6 +369,9 @@ class AdamW_adv(torch.optim.Optimizer):
         if self.use_AdEMAMix:
             beta3_ema = group['beta3_ema']
             alpha = group['alpha']
+        nesterov = group.get('nesterov', False)
+        nesterov_coef = group.get('nesterov_coef', None)
+        use_mt = group['betas'][0] > 0
 
         if group.get('kourkoutas_beta', False):
             # Accumulate current grad's norm for the *next* step
@@ -365,7 +384,7 @@ class AdamW_adv(torch.optim.Optimizer):
             grad_reshaped = grad.view(d1, d2)
 
             # Reconstruct momentum from previous step's factors
-            if beta1 > 0:
+            if use_mt:
                 mt = _reconstruct_state((state['mu_m_nmf'], state['mv_m_nmf'], state['sign'], d2), signed=True)
 
                 # Update momentum in full-size
@@ -381,6 +400,10 @@ class AdamW_adv(torch.optim.Optimizer):
                 else:
                     update_mt = mt
 
+                if nesterov:
+                    nv_coef = beta1 if nesterov_coef is None else nesterov_coef
+                    update_mt = update_mt.lerp_(grad_reshaped, 1-nv_coef)
+
             vt = _reconstruct_state((state['mu_v_nmf'], state['mv_v_nmf']), signed=False)
 
             if isinstance(beta2, torch.Tensor) and beta2.dim() > 0:
@@ -393,7 +416,7 @@ class AdamW_adv(torch.optim.Optimizer):
 
                 mt_slow.lerp_(grad_reshaped, 1.0 - beta3_ema)
 
-                if beta1 > 0:
+                if use_mt:
                     update = update_mt.add_(mt_slow, alpha=alpha)
                 else:
                     update = grad_reshaped.add(mt_slow, alpha=alpha)
@@ -402,7 +425,7 @@ class AdamW_adv(torch.optim.Optimizer):
                 state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'] = _factorize_state(mt_slow, signed=True)
                 del mt_slow
             else:
-                if beta1 > 0:
+                if use_mt:
                     update = update_mt
                 else:
                     update = grad_reshaped.clone()
@@ -429,7 +452,7 @@ class AdamW_adv(torch.optim.Optimizer):
             actual_precision = group['actual_state_precision']
             factored_2nd = state.get('factored_2nd', False)
 
-            if beta1 > 0:
+            if use_mt:
                 exp_avg = get_state(state, 'exp_avg', actual_precision)
                 exp_avg.lerp_(grad, 1.0 - beta1)
 
@@ -439,19 +462,24 @@ class AdamW_adv(torch.optim.Optimizer):
                     update_mt = _cautious_update(exp_avg, grad)
                 else:
                     update_mt = exp_avg.clone()
+
+                if nesterov:
+                    nv_coef = beta1 if nesterov_coef is None else nesterov_coef
+                    update_mt = update_mt.lerp_(grad, 1-nv_coef)
+
                 set_state(state, 'exp_avg', exp_avg, actual_precision, random_int_state_tensor)
 
             if self.use_AdEMAMix:
                 exp_avg_slow = get_state(state, 'exp_avg_slow', actual_precision)
                 exp_avg_slow.lerp_(grad, 1.0 - beta3_ema)
 
-                if beta1 > 0:
+                if use_mt:
                     update = update_mt.add_(exp_avg_slow, alpha=alpha)
                 else:
                     update = torch.add(grad, exp_avg_slow, alpha=alpha)
                 set_state(state, 'exp_avg_slow', exp_avg_slow, actual_precision, random_int_state_tensor)
             else:
-                update = update_mt if beta1 > 0 else grad.clone()
+                update = update_mt if use_mt else grad.clone()
 
             if factored_2nd:
                 d1, d2 = state['effective_shape']

{adv_optm-2.4.dev12 → adv_optm-2.4.dev14}/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, _init_fisher_wd_scaler, _get_fisher_wd_scaler
+from ..util.update_util import _grams_update, _cautious_update, _init_fisher_wd_scaler, _get_fisher_wd_scaler
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm, scale_eps
 from ..util.centered_decay import _init_anchor
 from ..util.state_util import init_state_tensor, get_state, set_state, upcast_grad_for_precision
@@ -32,7 +32,8 @@ class Adopt_adv(torch.optim.Optimizer):
         betas (tuple[float, float]): coefficients used for computing running
             averages of momentum and variance (default: (0.9, 0.9999))
         eps (float): term added to the denominator to improve
-            numerical stability (default: 1e-6)
+            numerical stability. Set to None for scale invariant eps (vector
+            lower bound) (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
@@ -68,16 +69,6 @@ class Adopt_adv(torch.optim.Optimizer):
             before it is added to the fast momentum term (`update = mt + alpha * mt_slow`).
             A higher value increases the stabilizing influence of the slow
             momentum. (default: 5.0)
-        Simplified_AdEMAMix (bool): whether to use the Simplified AdEMAMix update rule.
-            This changes the EMA to accumulator and the update numerator to `alpha_grad * grad + mt`, which can be
-            more responsive, especially for small batch sizes. Enabling this will
-            automatically disable `use_AdEMAMix`, `cautious_mask`, `grams_moment`,
-            and `use_atan2`. (default: False)
-        alpha_grad (float): Mixing coefficient for the Simplified AdEMAMix update rule
-            (only used when `Simplified_AdEMAMix` is `True`). Controls the weight of the
-            current gradient. For small batch sizes, use high values (e.g., 10-100) to be
-            more responsive. For large batch sizes, use low values (e.g., 0-1) for
-            stability. (default: 100.0)
         kourkoutas_beta (bool): whether to enable the layer-wise dynamic β₂ logic.
             If `False`, the optimizer behaves as standard Adopt. (default: False)
         beta2_min (float): The minimum value for dynamic β₂, used during periods of
@@ -143,9 +134,9 @@ class Adopt_adv(torch.optim.Optimizer):
         use_AdEMAMix: bool = False,
         beta3_ema: float = 0.9999,
         alpha: float = 5.0,
-        # One-EMA AdEMAMix
-        Simplified_AdEMAMix: bool = False,
-        alpha_grad: float = 100.0,
+        # Nesterov momentum
+        nesterov: bool = False,
+        nesterov_coef: float | None = None,
         # K-b (adaptive beta2)
         kourkoutas_beta: bool = False,
         beta2_min: float = 0.9,
@@ -179,16 +170,8 @@ class Adopt_adv(torch.optim.Optimizer):
         if cautious_mask and grams_moment:
             print("Warning: cautious is incompatible with grams, Disabling cautious.")
             cautious_mask = False
-        if betas[0] == 0.0 and Simplified_AdEMAMix:
-            raise ValueError(f"Beta1 cannot be 0.0 when using Simplified_AdEMAMix. Got {betas[0]}")
         if kourkoutas_beta and not (betas[1] > beta2_min):
             raise ValueError(f"For Kourkoutas-β, betas[1] (as beta2_max) must be > beta2_min. Got {betas[1]} and {beta2_min}")
-        if use_AdEMAMix and Simplified_AdEMAMix:
-            print("Warning: use_AdEMAMix is incompatible with Simplified_AdEMAMix, Disabling use_AdEMAMix.")
-        if grams_moment and Simplified_AdEMAMix:
-            print("Warning: grams is incompatible with Simplified_AdEMAMix, Disabling grams.")
-        if cautious_mask and Simplified_AdEMAMix:
-            print("Warning: cautious is incompatible with Simplified_AdEMAMix, Disabling cautious.")
 
 
         state_precision = state_precision.lower()
@@ -204,7 +187,7 @@ class Adopt_adv(torch.optim.Optimizer):
             "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,
+            "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,
             "spectral_normalization": spectral_normalization,
@@ -216,17 +199,18 @@ class Adopt_adv(torch.optim.Optimizer):
         }
         self.clip_lambda = clip_lambda
         self.stochastic_rounding = stochastic_rounding
-        self.use_atan2 = use_atan2 and not Simplified_AdEMAMix
-        self.cautious_mask = cautious_mask and not Simplified_AdEMAMix
-        self.grams_moment = grams_moment and not Simplified_AdEMAMix
+        self.use_atan2 = use_atan2
+        self.cautious_mask = cautious_mask
+        self.grams_moment = grams_moment
         self.orthogonal_gradient = orthogonal_gradient
-        self.use_AdEMAMix = use_AdEMAMix and not Simplified_AdEMAMix
-        self.Simplified_AdEMAMix = Simplified_AdEMAMix
+        self.use_AdEMAMix = use_AdEMAMix
         self.kourkoutas_beta = kourkoutas_beta
         self.layer_key_fn = layer_key_fn
         self._init_lr = lr
         super().__init__(params, defaults)
 
+        self.init_step()
+
         if self.kourkoutas_beta:
             self.kourkoutas_helper = KourkoutasHelper(self)
 
@@ -258,26 +242,15 @@ class Adopt_adv(torch.optim.Optimizer):
     @property
     def supports_flat_params(self): return False
 
-    @torch.no_grad()
-    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
-        if p.grad is None:
-            return
+    def init_step(self):
+        for group in self.param_groups:
+            for i, p in enumerate(group['params']):
+                self.__init_state(p, group)
 
-        grad = p.grad
+    @torch.no_grad()
+    def __init_state(self, p, group):
         state = self.state[p]
 
-        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
@@ -340,6 +313,27 @@ class Adopt_adv(torch.optim.Optimizer):
 
             _init_fisher_wd_scaler(group, state, p)
 
+    @torch.no_grad()
+    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
+        if p.grad is None:
+            return
+
+        grad = p.grad
+        state = self.state[p]
+        self.__init_state(p, group)
+
+        beta1, beta2 = group['betas']
+
+        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)
+
         current_step = state['step']
 
         # The first step is for initialization only (skip when use_atan2 as it's scale invariant).
@@ -367,9 +361,6 @@ class Adopt_adv(torch.optim.Optimizer):
             lr = group['lr']
             step_param_fn = self._step_parameter
 
-        if self.Simplified_AdEMAMix:
-            lr = _scale_sim_AdEMAMix_update(beta1, state['step'] + 1, group["alpha_grad"], lr, group.get('spectral_normalization', False))
-
         step_param_fn(p, grad, state, group, lr, beta1, beta2, random_int_tensor, random_int_state_tensor)
 
         state['step'] += 1
@@ -383,8 +374,9 @@ class Adopt_adv(torch.optim.Optimizer):
         if self.use_AdEMAMix:
             beta3_ema = group['beta3_ema']
             alpha = group['alpha']
-        if self.Simplified_AdEMAMix:
-            alpha_grad = group["alpha_grad"]
+        nesterov = group.get('nesterov', False)
+        nesterov_coef = group.get('nesterov_coef', None)
+        use_mt = group['betas'][0] > 0
 
         if group.get('kourkoutas_beta', False):
             # Accumulate current grad's norm for the *next* step
@@ -421,13 +413,10 @@ class Adopt_adv(torch.optim.Optimizer):
                     normalized_grad.clamp_(-clip_val, clip_val)
 
             # ADOPT Step B: Update momentum m_t using normalized gradient
-            if beta1 > 0:
+            if use_mt:
                 mt = _reconstruct_state((state['mu_m_nmf'], state['mv_m_nmf'], state['sign'], d2), signed=True)
 
-                if self.Simplified_AdEMAMix:
-                    mt.mul_(beta1).add_(normalized_grad, alpha=1.0)
-                else:
-                    mt.lerp_(normalized_grad, 1.0 - beta1)
+                mt.lerp_(normalized_grad, 1.0 - beta1)
 
                 # Factorize
                 state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(mt.clone(), signed=True)
@@ -439,13 +428,17 @@ class Adopt_adv(torch.optim.Optimizer):
                 else:
                     update_mt = mt
 
+                if nesterov:
+                    nv_coef = beta1 if nesterov_coef is None else nesterov_coef
+                    update_mt = update_mt.lerp_(grad_reshaped, 1-nv_coef)
+
             if self.use_AdEMAMix:
                 # Reconstruct AdEMAMix EMA
                 mt_slow = _reconstruct_state((state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'], d2), signed=True)
 
                 mt_slow.lerp_(normalized_grad, 1.0 - beta3_ema)
 
-                if beta1 > 0:
+                if use_mt:
                     update = update_mt.add_(mt_slow, alpha=alpha)
                     del normalized_grad
                 else:
@@ -453,12 +446,8 @@ class Adopt_adv(torch.optim.Optimizer):
                 # Factorize
                 state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'] = _factorize_state(mt_slow, signed=True)
                 del mt_slow
-
-            elif self.Simplified_AdEMAMix:
-                update = update_mt.add_(normalized_grad, alpha=alpha_grad)
-                del normalized_grad
             else:
-                if beta1 > 0:
+                if use_mt:
                     update = update_mt
                     del normalized_grad
                 else:
@@ -490,12 +479,9 @@ class Adopt_adv(torch.optim.Optimizer):
                     normalized_grad.clamp_(-clip_val, clip_val)
 
             # ADOPT Step B: Update momentum m_t
-            if beta1 > 0:
+            if use_mt:
                 mt = get_state(state, 'exp_avg', actual_precision) # m_{t-1}
-                if self.Simplified_AdEMAMix:
-                    mt.mul_(beta1).add_(normalized_grad, alpha=1.0)
-                else:
-                    mt.lerp_(normalized_grad, 1.0 - beta1)
+                mt.lerp_(normalized_grad, 1.0 - beta1)
 
                 if self.grams_moment:
                     update_mt = _grams_update(mt, grad)
@@ -504,21 +490,23 @@ class Adopt_adv(torch.optim.Optimizer):
                 else:
                     update_mt = mt.clone()
 
+                if nesterov:
+                    nv_coef = beta1 if nesterov_coef is None else nesterov_coef
+                    update_mt = update_mt.lerp_(grad, 1-nv_coef)
+
                 set_state(state, 'exp_avg', mt, actual_precision, random_int_state_tensor)
 
             if self.use_AdEMAMix:
                 m_slow = get_state(state, 'exp_avg_slow', actual_precision)
                 m_slow.lerp_(normalized_grad, 1.0 - beta3_ema)
-                if beta1 > 0:
+                if use_mt:
                     update = update_mt.add_(m_slow, alpha=alpha)
                     del normalized_grad
                 else:
                     update = normalized_grad.add_(m_slow, alpha=alpha)
                 set_state(state, 'exp_avg_slow', m_slow, actual_precision, random_int_state_tensor)
-            elif self.Simplified_AdEMAMix:
-                update = update_mt.add_(normalized_grad, alpha=alpha_grad)
             else:
-                if beta1 > 0:
+                if use_mt:
                     update = update_mt
                     del normalized_grad
                 else: