adv-optm 0.1.6__tar.gz → 0.1.8__tar.gz

{adv_optm-0.1.6 → adv_optm-0.1.8}/PKG-INFO

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

{adv_optm-0.1.6 → adv_optm-0.1.8}/adv_optm/__init__.py

@@ -2,6 +2,7 @@ from .optim import (
     AdamW_adv,
     Prodigy_adv,
     Adopt_adv,
+    Simplified_AdEMAMix,
     Lion_adv,
     Lion_Prodigy_adv,
 )
@@ -10,8 +11,9 @@ __all__ = [
     "AdamW_adv",
     "Prodigy_adv",
     "Adopt_adv",
+    "Simplified_AdEMAMix",
     "Lion_adv",
     "Lion_Prodigy_adv",
 ]
 
-__version__ = "0.1.6"
+__version__ = "0.1.8"

{adv_optm-0.1.6 → adv_optm-0.1.8}/adv_optm/optim/AdamW_adv.py

@@ -55,7 +55,7 @@ class AdamW_adv(torch.optim.Optimizer):
             the warmup, `alpha` ramps from 0 to its target value. If `None`,
             the scheduler is disabled. (default: None)
         factored (bool): whether to use the factorization or disable it to use
-            the uncompressed optimizer. (default: True)
+            the uncompressed optimizer. (default: False)
     """
 
     def __init__(
@@ -76,7 +76,7 @@ class AdamW_adv(torch.optim.Optimizer):
         beta3_ema: float = 0.9999,
         alpha: float = 5.0,
         t_alpha: int | None = None,
-        factored: bool = True,
+        factored: bool = False,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -216,7 +216,10 @@ class AdamW_adv(torch.optim.Optimizer):
                 del unpacked_sign_slow
 
                 mt_slow.mul_(beta3_ema).add_(grad_reshaped, alpha=1.0 - beta3_ema)
-                update = mt + (alpha_t * mt_slow) if beta1 > 0 else grad_reshaped + (alpha_t * mt_slow)
+                if beta1 > 0:
+                    update = torch.add(mt, mt_slow, alpha=alpha_t)
+                else:
+                    update = torch.add(grad_reshaped, mt_slow, alpha=alpha_t)
             else:
                 update = mt.clone() if beta1 > 0 else grad_reshaped.clone()
             del grad_reshaped
@@ -262,7 +265,10 @@ class AdamW_adv(torch.optim.Optimizer):
             if self.use_AdEMAMix:
                 exp_avg_slow = state['exp_avg_slow']
                 exp_avg_slow.mul_(beta3_ema).add_(grad, alpha=1 - beta3_ema)
-                update = exp_avg + (alpha_t * exp_avg_slow) if beta1 > 0 else grad + (alpha_t * exp_avg_slow)
+                if beta1 > 0:
+                    update = torch.add(exp_avg, exp_avg_slow, alpha=alpha_t)
+                else:
+                    update = torch.add(grad, exp_avg_slow, alpha=alpha_t)
             else:
                 update = exp_avg.clone() if beta1 > 0 else grad.clone()
 

{adv_optm-0.1.6 → adv_optm-0.1.8}/adv_optm/optim/Adopt_adv.py

@@ -63,7 +63,7 @@ class Adopt_adv(torch.optim.Optimizer):
             the scheduler is disabled and the full `alpha` value is used from
             the start. (default: None)
         factored (bool): whether to use the factorization or disable it to use
-            the uncompressed optimizer. (default: True)
+            the uncompressed optimizer. (default: False)
     """
 
     def __init__(
@@ -84,7 +84,7 @@ class Adopt_adv(torch.optim.Optimizer):
         beta3_ema: float = 0.9999,
         alpha: float = 5.0,
         t_alpha: int | None = None,
-        factored: bool = True,
+        factored: bool = False,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -235,7 +235,7 @@ class Adopt_adv(torch.optim.Optimizer):
 
             if self.use_AdEMAMix:
                 mt_slow.mul_(beta3_ema).add_(normalized_grad, alpha=1.0 - beta3_ema)
-                update = mt + (alpha_t * mt_slow)
+                update = torch.add(mt, m_slow, alpha=alpha_t)
                 update = update.view(p.shape)
             else:
                 update = mt.view(p.shape)
@@ -295,9 +295,9 @@ class Adopt_adv(torch.optim.Optimizer):
 
             if self.use_AdEMAMix:
                 m_slow.mul_(beta3_ema).add_(normalized_grad, alpha=1.0 - beta3_ema)
-                update = m + (alpha_t * m_slow)
+                update = torch.add(m, m_slow, alpha=alpha_t)
             else:
-                update = m
+                update = m.clone()
 
             if self.use_atan2:
                 update.mul_(group['lr'] * 1.2732395447351628)

{adv_optm-0.1.6 → adv_optm-0.1.8}/adv_optm/optim/Lion_Prodigy_adv.py

@@ -33,8 +33,6 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
             (default: 0.0).
         factored (bool): whether to use the factorization or use the
             uncompressed optimizer. (default: True)
-        variance_reduction (bool): whether to use the variance reduction technique
-            from "Convergence Analysis of the Lion Optimizer" (arXiv:2508.12327v1). (default: False).
         d0 (float):
             Initial D estimate for D-adaptation (default 1e-6). Rarely needs changing.
         d_coef (float):
@@ -66,7 +64,6 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
         use_cautious: bool = False,
         clip_threshold: float = 0.0,
         factored: bool = True,
-        variance_reduction: bool = False,
         # prodigy parameters
         beta3: float = None,
         d0: float = 1e-6,
@@ -97,7 +94,6 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
         self.stochastic_rounding = stochastic_rounding
         self.use_cautious = use_cautious
         self.factored = factored
-        self.variance_reduction = variance_reduction
         self.fsdp_in_use = fsdp_in_use
         super().__init__(params, defaults)
         # Global state for accumulating metrics across parameter updates within a single step.
@@ -183,12 +179,8 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
                 state['mv_m_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
                 packed_d2 = (d2 + 7) // 8
                 state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
-                if self.variance_reduction:
-                    state['prev_grad'] = torch.zeros((d1, d2), device=p.device, dtype=dtype)
             else: # Fallback to standard Lion
                 state['exp_avg'] = torch.zeros_like(p, device=p.device, dtype=dtype)
-                if self.variance_reduction:
-                    state['prev_grad'] = torch.zeros_like(p, device=p.device, dtype=dtype)
 
         if state['factored']:
             # Factored Path
@@ -215,20 +207,7 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
             update_for_param = signed_update.view(p.shape).mul(self.dlr)
 
             # Update momentum m_t = β2*m_{t-1} + (1-β2)*lr*g_t
-            if self.variance_reduction:
-                if state['step'] == 1:
-                    exp_avg.copy_(grad_reshaped)
-                else:
-                    # Heuristic Prodigy-STORM update
-                    correction = exp_avg.sub(state['prev_grad'])
-                    grad_alpha = self.d * (1 - self.beta2) + self.beta2
-                    exp_avg.copy_(grad_reshaped).mul_(grad_alpha).add_(correction, alpha=self.beta2)
-                    del correction, grad_alpha
-                state['prev_grad'].copy_(grad_reshaped)
-            else:
-                # Standard Prodigy-Lion
-                alpha = self.d * (1 - self.beta2)
-                exp_avg.mul_(self.beta2).add_(grad_reshaped, alpha=alpha)
+            exp_avg.mul_(self.beta2).add_(grad_reshaped, alpha=self.d * (1 - self.beta2))
             del grad_reshaped
 
             # Compress new momentum m_t and store factors
@@ -254,20 +233,7 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
             update_for_param = signed_update.mul(self.dlr)
 
             # Update momentum 
-            if self.variance_reduction:
-                if state['step'] == 1:
-                    exp_avg.copy_(grad)
-                else:
-                    # Heuristic Prodigy-STORM update
-                    correction = exp_avg.sub(state['prev_grad'])
-                    grad_alpha = self.d * (1 - self.beta2) + self.beta2
-                    exp_avg.copy_(grad).mul_(grad_alpha).add_(correction, alpha=self.beta2)
-                    del grad_alpha, correction
-                state['prev_grad'].copy_(grad)
-            else:
-                # Standard Prodigy-Lion
-                alpha = self.d * (1 - self.beta2)
-                exp_avg.mul_(self.beta2).add_(grad, alpha=alpha)
+            exp_avg.mul_(self.beta2).add_(grad, alpha=self.d * (1 - self.beta2))
 
         # --- Accumulate Prodigy stats ---
         d0, safeguard_warmup, slice_p = group['d0'], group['safeguard_warmup'], group['slice_p']
@@ -298,7 +264,7 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
         else:
             p.data.add_(-update_for_param)
 
-            del update_for_param
+        del update_for_param
 
     @torch.no_grad()
     def step(self, closure: Optional[callable] = None):

{adv_optm-0.1.6 → adv_optm-0.1.8}/adv_optm/optim/Lion_adv.py

@@ -33,8 +33,6 @@ class Lion_adv(torch.optim.Optimizer):
             (default: 0.0).
         factored (bool): whether to use the factorization or use the
             uncompressed optimizer. (default: True)
-        variance_reduction (bool): whether to use the variance reduction technique
-            from "Convergence Analysis of the Lion Optimizer" (arXiv:2508.12327v1). (default: False).
     """
 
     def __init__(
@@ -49,7 +47,6 @@ class Lion_adv(torch.optim.Optimizer):
         use_cautious: bool = False,
         clip_threshold: float = 0.0,
         factored: bool = True,
-        variance_reduction: bool = False,
     ):
         if not lr > 0.0:
             raise ValueError(f"Learning rate must be > 0.0, but got {lr}")
@@ -69,7 +66,6 @@ class Lion_adv(torch.optim.Optimizer):
         self.stochastic_rounding = stochastic_rounding
         self.use_cautious = use_cautious
         self.factored = factored
-        self.variance_reduction = variance_reduction
         super().__init__(params, defaults)
 
     @property
@@ -122,12 +118,8 @@ class Lion_adv(torch.optim.Optimizer):
                 state['mv_m_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
                 packed_d2 = (d2 + 7) // 8
                 state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
-                if self.variance_reduction:
-                    state['prev_grad'] = torch.zeros((d1, d2), device=p.device, dtype=dtype)
             else: # Fallback to standard Lion
                 state['exp_avg'] = torch.zeros_like(p, device=p.device, dtype=dtype)
-                if self.variance_reduction:
-                    state['prev_grad'] = torch.zeros_like(p, device=p.device, dtype=dtype)
 
         state['step'] += 1
         beta1, beta2 = group["betas"]
@@ -157,21 +149,9 @@ class Lion_adv(torch.optim.Optimizer):
             # Parameter update
             update_for_param = signed_update.view(p.shape).mul_(lr)
 
-            # Update momentum
-            if self.variance_reduction:
-                if state['step'] == 1:
-                    exp_avg.copy_(grad_reshaped)
-                else:
-                    # Use the simplified STORM update: m_t = g_t + β₂ * (m_{t-1} - g_{t-1})
-                    correction = exp_avg.sub(state['prev_grad'])
-                    # Calculate the new momentum and store it back into exp_avg
-                    exp_avg.copy_(grad_reshaped).add_(correction, alpha=beta2)
-                    del correction
-                # Update prev_grad for the next iteration
-                state['prev_grad'].copy_(grad_reshaped)
-            else:
-                # Standard Lion momentum update
-                exp_avg.mul_(beta2).add_(grad_reshaped, alpha=1-beta2)
+            # Standard Lion momentum update
+            exp_avg.mul_(beta2).add_(grad_reshaped, alpha=1-beta2)
+            del grad_reshaped
 
             # Compress new momentum m_t and store factors
             state['sign'] = _pack_bools(exp_avg > 0)
@@ -195,21 +175,8 @@ class Lion_adv(torch.optim.Optimizer):
 
             update_for_param = signed_update.mul_(lr)
 
-            # Update momentum 
-            if self.variance_reduction:
-                if state['step'] == 1:
-                    exp_avg.copy_(grad)
-                else:
-                    # Use the simplified STORM update: m_t = g_t + β₂ * (m_{t-1} - g_{t-1})
-                    correction = exp_avg.sub(state['prev_grad'])
-                    # Calculate the new momentum and store it back into exp_avg
-                    exp_avg.copy_(grad).add_(correction, alpha=beta2)
-                    del correction
-                # Update prev_grad for the next iteration
-                state['prev_grad'].copy_(grad)
-            else:
-                # Standard Lion momentum update
-                exp_avg.mul_(beta2).add_(grad, alpha=1-beta2)
+            # Standard Lion momentum update
+            exp_avg.mul_(beta2).add_(grad, alpha=1-beta2)
 
         if group["weight_decay"] != 0:
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
@@ -225,7 +192,7 @@ class Lion_adv(torch.optim.Optimizer):
         else:
             p.data.add_(-update_for_param)
 
-            del update_for_param
+        del update_for_param
 
     @torch.no_grad()
     def step(self, closure: Optional[callable] = None):

{adv_optm-0.1.6 → adv_optm-0.1.8}/adv_optm/optim/Prodigy_adv.py

@@ -52,9 +52,19 @@ class Prodigy_adv(torch.optim.Optimizer):
             highly recommended to prevent instability at the beginning of training,
             as it gradually introduces the stabilizing slow momentum term. During
             the warmup, `alpha` ramps from 0 to its target value. If `None`,
-            the scheduler is disabled and th
+            the scheduler is disabled.
+        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`, `use_cautious`, `use_grams`,
+            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)
         factored (bool): whether to use the factorization or disable it to use
-            the uncompressed optimizer. (default: True)
+            the uncompressed optimizer. (default: False)
         d0 (float):
             Initial D estimate for D-adaptation (default 1e-6). Rarely needs changing.
         d_coef (float):
@@ -72,6 +82,9 @@ class Prodigy_adv(torch.optim.Optimizer):
         slice_p (int): Reduce memory usage by calculating LR adaptation statistics on only every 
             pth entry of each tensor. For values greater than 1 this an an approximation to standard 
             Prodigy. Values ~11 are reasonable (default 11).
+        prodigy_steps (int): If greater than zero, disable Prodigy's stepsize adjustments
+            after the specified optimiser step and release all state memory required by Prodigy
+            (default: 0).
     """
 
     def __init__(
@@ -91,7 +104,9 @@ class Prodigy_adv(torch.optim.Optimizer):
         beta3_ema: float = 0.9999,
         alpha: float = 5.0,
         t_alpha: int | None = None,
-        factored: bool = True,
+        Simplified_AdEMAMix: bool = False,
+        alpha_grad: float = 100.0,
+        factored: bool = False,
         # prodigy parameters
         beta3: float = None,
         d0: float = 1e-6,
@@ -100,6 +115,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         safeguard_warmup: bool = False,
         fsdp_in_use: bool = False,
         slice_p: int = 11,
+        prodigy_steps: int = 0,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -109,6 +125,22 @@ class Prodigy_adv(torch.optim.Optimizer):
             raise ValueError(f"Epsilon should be >= 0.0. Got {eps}")
         if not (weight_decay >= 0.0):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
+        if not (prodigy_steps >= 0):
+            raise ValueError(f"prodigy_steps should be >= 0. Got {prodigy_steps}")
+        if betas[0] == 0.0 and Simplified_AdEMAMix:
+            raise ValueError(f"Beta 1 cannot be 0.0 when using Simplified_AdEMAMix. Got {betas[0]}")
+        if use_AdEMAMix and Simplified_AdEMAMix:
+            print("Warning: use_AdEMAMix is incompatible with Simplified_AdEMAMix, Disabling use_AdEMAMix.")
+        if use_grams and Simplified_AdEMAMix:
+            print("Warning: use_grams is incompatible with Simplified_AdEMAMix, Disabling use_grams.")
+        if use_cautious and Simplified_AdEMAMix:
+            print("Warning: use_cautious is incompatible with Simplified_AdEMAMix, Disabling use_cautious.")
+        if use_atan2 and Simplified_AdEMAMix:
+            print("Warning: use_atan2 is incompatible with Simplified_AdEMAMix. Disabling use_atan2.")
+            use_atan2 = False
+        if Simplified_AdEMAMix and alpha_grad > 0:
+            # scales d_coef by alpha_grad, this force prodigy to behave well with Simplified_AdEMAMix
+            d_coef = d_coef/alpha_grad
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
@@ -117,12 +149,14 @@ class Prodigy_adv(torch.optim.Optimizer):
             "beta3_ema": beta3_ema, "alpha": alpha, "t_alpha": t_alpha,
             "beta3": beta3, "d": d0, "d0": d0, "d_max": d0, "d_numerator": 0.0, "d_coef": d_coef,
             "growth_rate": growth_rate, "safeguard_warmup": safeguard_warmup, "k": 0, "slice_p": slice_p,
-            "fsdp_in_use": fsdp_in_use,
+            "fsdp_in_use": fsdp_in_use, "prodigy_steps": prodigy_steps,
+            "alpha_grad": alpha_grad, 
         }
         self.stochastic_rounding = stochastic_rounding
-        self.use_cautious = use_cautious
-        self.use_grams = use_grams
-        self.use_AdEMAMix = use_AdEMAMix
+        self.use_cautious = use_cautious and not Simplified_AdEMAMix
+        self.use_grams = use_grams and not Simplified_AdEMAMix
+        self.use_AdEMAMix = use_AdEMAMix and not Simplified_AdEMAMix
+        self.Simplified_AdEMAMix = Simplified_AdEMAMix
         self.factored = factored
         self.fsdp_in_use = fsdp_in_use
         super().__init__(params, defaults)
@@ -229,6 +263,8 @@ class Prodigy_adv(torch.optim.Optimizer):
             alpha_t = alpha
             if t_alpha is not None and t_alpha > 0 and current_step < t_alpha:
                 alpha_t = min(current_step * alpha / t_alpha, alpha)
+        if self.Simplified_AdEMAMix:
+            alpha_grad = group["alpha_grad"]
 
         if state['factored']:
             d1, d2 = state['effective_shape']
@@ -243,7 +279,10 @@ class Prodigy_adv(torch.optim.Optimizer):
                     torch.where(unpacked_sign, mt, -mt, out=mt)
                     del unpacked_sign
                 # Update momentum in full-size
-                mt.mul_(self.beta1).add_(grad_reshaped, alpha=self.d * (1.0 - self.beta1))
+                if self.Simplified_AdEMAMix:
+                    mt.mul_(self.beta1).add_(grad_reshaped, alpha=self.d)
+                else:
+                    mt.mul_(self.beta1).add_(grad_reshaped, alpha=self.d * (1.0 - self.beta1))
                 if self.use_grams:
                     mt.copy_(grad_reshaped.sign() * mt.abs())
                 elif self.use_cautious:
@@ -263,7 +302,12 @@ class Prodigy_adv(torch.optim.Optimizer):
                 torch.where(unpacked_sign_slow, mt_slow, -mt_slow, out=mt_slow)
                 del unpacked_sign_slow
                 mt_slow.mul_(beta3_ema).add_(grad_reshaped, alpha=self.d * (1.0 - beta3_ema))
-                update = mt + (alpha_t * mt_slow) if self.beta1 > 0 else grad_reshaped + (alpha_t * mt_slow)
+                if self.beta1 > 0:
+                    update = torch.add(mt, mt_slow, alpha=alpha_t)
+                else:
+                    update = torch.add(grad_reshaped, mt_slow, alpha=alpha_t)
+            elif self.Simplified_AdEMAMix:
+                update = torch.add(mt, grad_reshaped, alpha=alpha_grad * self.d)
             else:
                 update = mt.clone() if self.beta1 > 0 else grad_reshaped.clone()
             del grad_reshaped
@@ -297,7 +341,10 @@ class Prodigy_adv(torch.optim.Optimizer):
 
             if self.beta1 > 0:
                 exp_avg = state['exp_avg']
-                exp_avg.mul_(self.beta1).add_(grad, alpha=self.d * (1.0 - self.beta1))
+                if self.Simplified_AdEMAMix:
+                    exp_avg.mul_(self.beta1).add_(grad, alpha=self.d)
+                else:
+                    exp_avg.mul_(self.beta1).add_(grad, alpha=self.d * (1.0 - self.beta1))
                 if self.use_grams:
                     exp_avg = grad.sign() * exp_avg.abs()
                 elif self.use_cautious:
@@ -309,7 +356,12 @@ class Prodigy_adv(torch.optim.Optimizer):
             if self.use_AdEMAMix:
                 exp_avg_slow = state['exp_avg_slow']
                 exp_avg_slow.mul_(beta3_ema).add_(grad, alpha=self.d * (1.0 - beta3_ema))
-                update = exp_avg + (alpha_t * exp_avg_slow) if self.beta1 > 0 else grad + (alpha_t * exp_avg_slow)
+                if self.beta1 > 0:
+                    update = torch.add(exp_avg, exp_avg_slow, alpha=alpha_t)
+                else:
+                    update = torch.add(grad, exp_avg_slow, alpha=alpha_t)
+            elif self.Simplified_AdEMAMix:
+                update = torch.add(exp_avg, grad, alpha=alpha_grad * self.d)
             else:
                 update = exp_avg.clone() if self.beta1 > 0 else grad.clone()
 
@@ -327,19 +379,27 @@ class Prodigy_adv(torch.optim.Optimizer):
             update.mul_(self.dlr)
 
         # --- Accumulate Prodigy stats ---
-        d0, safeguard_warmup, slice_p = group['d0'], group['safeguard_warmup'], group['slice_p']
-        s, p0 = state['s'], state['p0']
-        grad_flat = grad.flatten().float()
-        p_flat = p.data.flatten().float()
-        p0 = p0.float()
+        prodigy_steps = group['prodigy_steps']
+        if prodigy_steps <= 0 or group['k'] < prodigy_steps:
+            d0, safeguard_warmup, slice_p = group['d0'], group['safeguard_warmup'], group['slice_p']
+            s, p0 = state['s'], state['p0']
+            grad_flat = grad.flatten().float()
+            p_flat = p.data.flatten().float()
+            p0 = p0.float()
 
-        self.d_numerator += (self.d / d0) * self.dlr * torch.dot(grad_flat[::slice_p], p0.data - p_flat[::slice_p]).item()
+            self.d_numerator += (self.d / d0) * self.dlr * torch.dot(grad_flat[::slice_p], p0.data - p_flat[::slice_p]).item()
 
-        alpha = ((self.d / d0) * self.d) if safeguard_warmup else ((self.d / d0) * self.dlr)
-        s.mul_(self.beta3).add_(grad_flat[::slice_p], alpha=alpha)
-        self.d_denom += s.abs().sum().item()
+            alpha = ((self.d / d0) * self.d) if safeguard_warmup else ((self.d / d0) * self.dlr)
+            s.mul_(self.beta3).add_(grad_flat[::slice_p], alpha=alpha)
+            self.d_denom += s.abs().sum().item()
 
-        del s, p0, grad_flat, p_flat, alpha
+            del s, p0, grad_flat, p_flat, alpha
+        else:
+            # Free memory if prodigy_steps is reached
+            if 's' in state:
+                del state['s']
+            if 'p0' in state:
+                del state['p0']
 
         # Decoupled weight decay
         if group["weight_decay"] != 0:
@@ -376,29 +436,37 @@ class Prodigy_adv(torch.optim.Optimizer):
     def calculate_d(self):
         """Calculates the new `d` based on the accumulated stats."""
         g_group = self.param_groups[0]
-        d_max, d_coef, growth_rate = g_group['d_max'], g_group['d_coef'], g_group['growth_rate']
         
-        if self.fsdp_in_use and dist.is_available() and dist.is_initialized():
-            # Use the device of the first parameter to avoid hardcoding '.cuda()'
-            device = self.param_groups[0]['params'][0].device
-            dist_tensor = torch.tensor([self.d_numerator, self.d_denom], device=device)
-            dist.all_reduce(dist_tensor, op=dist.ReduceOp.SUM)
-            global_d_numerator = dist_tensor[0].item()
-            global_d_denom = dist_tensor[1].item()
-        else:
-            global_d_numerator = self.d_numerator
-            global_d_denom = self.d_denom
-
-        d_hat = self.d
-        if global_d_denom > 0:
-            d_hat = d_coef * global_d_numerator / global_d_denom
-            if self.d == g_group['d0']:
-                self.d = max(self.d, d_hat)
-            d_max = max(d_max, d_hat)
-            self.d = min(d_max, self.d * growth_rate)
-
+        # Only perform d-adaptation if prodigy_steps has not been reached
+        prodigy_active = not (g_group.get('prodigy_steps', 0) > 0 and g_group['k'] >= g_group['prodigy_steps'])
+
+        if prodigy_active:
+            d_max, d_coef, growth_rate = g_group['d_max'], g_group['d_coef'], g_group['growth_rate']
+            
+            if self.fsdp_in_use and dist.is_available() and dist.is_initialized():
+                # Use the device of the first parameter to avoid hardcoding '.cuda()'
+                device = self.param_groups[0]['params'][0].device
+                dist_tensor = torch.tensor([self.d_numerator, self.d_denom], device=device)
+                dist.all_reduce(dist_tensor, op=dist.ReduceOp.SUM)
+                global_d_numerator = dist_tensor[0].item()
+                global_d_denom = dist_tensor[1].item()
+            else:
+                global_d_numerator = self.d_numerator
+                global_d_denom = self.d_denom
+
+            d_hat = self.d
+            if global_d_denom > 0:
+                d_hat = d_coef * global_d_numerator / global_d_denom
+                if self.d == g_group['d0']:
+                    self.d = max(self.d, d_hat)
+                d_max = max(d_max, d_hat)
+                self.d = min(d_max, self.d * growth_rate)
+
+            for group in self.param_groups:
+                group['d_numerator'] = global_d_numerator
+                group['d'] = self.d
+                group['d_max'] = d_max
+        
+        # Increment step counter for all groups, regardless of whether d was updated
         for group in self.param_groups:
-            group['d_numerator'] = global_d_numerator
-            group['d'] = self.d
-            group['d_max'] = d_max
             group['k'] += 1

adv_optm-0.1.8/adv_optm/optim/Simplified_AdEMAMix.py

@@ -0,0 +1,246 @@
+import torch
+
+import math
+
+from ..util.BF16_Stochastic_Rounding import add_stochastic_
+from ..util.Effective_Shape import _get_effective_shape
+from ..util.NNMF import _nnmf,_unnmf
+from ..util.OrthoGrad import _orthogonalize_gradient
+from ..util.One_Bit_Boolean import _pack_bools, _unpack_bools
+
+# A little helper from the original simplified_AdEMAMix
+def linear_hl_warmup_scheduler(step, beta_end, beta_start=0, warmup=1):
+
+    def f(beta, eps=1e-8):
+        return math.log(0.5)/math.log(beta+eps)-1
+
+    def f_inv(t):
+        return math.pow(0.5, 1/(t+1))
+
+    if step < warmup:
+        a = step / float(warmup)
+        return f_inv((1.0-a) * f(beta_start) + a * f(beta_end))
+    return beta_end
+
+class Simplified_AdEMAMix(torch.optim.Optimizer):
+    """
+    Implements the Simplified AdEMAMix algorithm.
+    Refactored from:
+    https://github.com/DepenM/Simplified-AdEMAMix/blob/main/simplified_AdEMAMix.py
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float): learning rate (default: 1e-5)
+        betas (tuple[float, float]): coefficients used for computing running
+            averages of gradient and its square (default: (0.99, 0.999))
+        eps (float): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float): weight decay (L2 penalty) (default: 0).
+        alpha_grad (float): Coeficient for mixing the current gradient and EMA. for small batch
+            sizes set it to high values, up to 100. And for large batch sized set it to small
+            value, down to 0. (default: 100)
+        beta1_warmup (int, optional): number of warmup steps used to increase beta1 (default: None)
+        min_beta1 (float, optional): minimum value of beta1 to start from (default 0.9)
+        vector_reshape (bool): whether to reshape 1D vectors into 2D
+            matrices to apply low-rank compression (default: True).
+        stochastic_rounding (bool): whether to use stochastic
+            rounding for BF16 parameter updates (default: True).
+        use_orthograd (bool): whether to use OrthoGrad. (default: False)
+        factored (bool): whether to use the factorization or disable it to use
+            the uncompressed optimizer. (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-5,
+        betas: tuple[float, float] = (0.99, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+        alpha_grad: float = 100.0,
+        beta1_warmup: int | None = None,
+        min_beta1: float | None = 0.9,
+        use_bias_correction: bool = True,
+        vector_reshape: bool = True,
+        stochastic_rounding: bool = True,
+        use_orthograd: bool = False,
+        factored: bool = False,
+    ):
+        if not (lr >= 0.0):
+            raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
+        if not (0.0 <= betas[0] < 1.0 and 0.0 <= betas[1] < 1.0):
+            raise ValueError(f"Betas should be in [0.0, 1.0). Got {betas}")
+        if not (eps >= 0.0):
+            raise ValueError(f"Epsilon should be >= 0.0. Got {eps}")
+        if not (weight_decay >= 0.0):
+            raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
+        if not 0.0 <= alpha_grad:
+            raise ValueError("Invalid alpha value: {}".format(alpha_grad))
+
+        defaults = {
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
+            "alpha_grad": alpha_grad, "beta1_warmup": beta1_warmup, "min_beta1": min_beta1,
+            "vector_reshape": vector_reshape,
+            "use_orthograd": use_orthograd, "use_bias_correction": use_bias_correction,
+        }
+        self.stochastic_rounding = stochastic_rounding
+        self.factored = factored
+        super().__init__(params, defaults)
+
+    @property
+    def supports_fused_back_pass(self):
+        return True
+
+    @property
+    def supports_memory_efficient_fp16(self):
+        return True
+
+    @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
+
+        grad = p.grad
+        if grad.dtype != torch.float32 and self.factored:
+            grad = grad.float()
+        if group["use_orthograd"]:
+            grad = _orthogonalize_gradient(p, grad)
+        state = self.state[p]
+
+        # State Initialization
+        if len(state) == 0:
+            state['step'] = 0
+
+            should_factor = (
+                self.factored and
+                not (len(p.shape) == 1 and not group['vector_reshape'])
+            )
+
+            state['factored'] = should_factor
+
+            dtype = torch.float32 if self.factored else p.dtype
+            device = p.device
+
+            if state['factored']:
+                state['effective_shape'] = _get_effective_shape(p.numel())
+                d1, d2 = state['effective_shape']
+
+                # First moment (m)
+                state['mu_m_nmf'] = torch.zeros(d1, device=device, dtype=dtype) 
+                state['mv_m_nmf'] = torch.zeros(d2, device=device, dtype=dtype)
+                packed_d2 = (d2 + 7) // 8
+                state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=device)
+                # Second moment (v)
+                state['mu_v_nmf'] = torch.zeros(d1, device=device, dtype=dtype) 
+                state['mv_v_nmf'] = torch.zeros(d2, device=device, dtype=dtype)
+            else:  # Fallback to standard optimizer for non-factored tensors
+                state['exp_avg'] = torch.zeros_like(p, device=device, dtype=dtype)
+                state['exp_avg_sq'] = torch.zeros_like(p, device=device, dtype=dtype)
+            
+            if group['use_bias_correction']:
+                state['num_sum'] = 0.0
+                state['den_sum'] = 0.0
+            else:
+                state['num_sum'] = 1.0
+                state['den_sum'] = 1.0
+
+        beta1_final, beta2 = group["betas"]
+        beta1_warmup = group["beta1_warmup"]
+        alpha_grad = group["alpha_grad"]
+
+        if beta1_warmup is not None:
+            step = state['step'] + 1
+            beta1 = linear_hl_warmup_scheduler(step, beta_end=beta1_final, beta_start=group['min_beta1'], warmup=beta1_warmup)
+        else:
+            beta1 = beta1_final
+
+        if group['use_bias_correction']:
+            state['num_sum'] = beta1 * state['num_sum'] + 1.0
+            state['den_sum'] = beta2 * state['den_sum'] + (1.0 - beta2)
+
+        if state['factored']:
+            d1, d2 = state['effective_shape']
+
+            # Reconstruct momentum from previous step's factors
+            mt = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
+            unpacked_sign = _unpack_bools(state['sign'], original_m=d2)
+            torch.where(unpacked_sign, mt, -mt, out=mt)
+            del unpacked_sign
+            # Update momentum in full-size
+            grad_reshaped = grad.view(d1, d2)
+            mt.mul_(beta1).add_(grad_reshaped, alpha=1.0)
+
+            vt = _unnmf((state['mu_v_nmf'], state['mv_v_nmf']))
+            vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
+
+            update = torch.add(mt, grad_reshaped, alpha=alpha_grad)
+            del grad_reshaped
+
+            denom = vt.sqrt().add_(group['eps'] * math.sqrt(state['den_sum']))
+            update.div_(denom)
+            del denom
+
+            if group['use_bias_correction']:
+                update = (update / state['num_sum']) * math.sqrt(state['den_sum'])
+
+            update = update.view(p.shape).mul_(group['lr'])
+
+            # Compress updated moments and store new factors
+            state['sign'] = _pack_bools(mt > 0)
+            _nnmf(mt.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
+            del mt
+            _nnmf(vt, out=(state['mu_v_nmf'], state['mv_v_nmf']))
+            del vt
+
+        else:  # Standard optimizer logic for non-factored tensors
+            exp_avg_sq = state['exp_avg_sq']
+
+            exp_avg = state['exp_avg']
+            exp_avg.mul_(beta1).add_(grad, alpha=1.0)
+
+            update = torch.add(exp_avg, grad, alpha=alpha_grad)
+
+            exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+            denom = exp_avg_sq.sqrt().add_(group['eps'] * math.sqrt(state['den_sum']))
+            update.div_(denom)
+            del denom
+
+            if group['use_bias_correction']:
+                update = (update / state['num_sum']) * math.sqrt(state['den_sum'])
+
+            update.mul_(group['lr'])
+
+        # Decoupled weight decay
+        if group["weight_decay"] != 0:
+            if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+                add_stochastic_(p.data, p.data, alpha=-group["weight_decay"] * group["lr"])
+            else:
+                p.data.add_(p.data, alpha=-group["weight_decay"] * group["lr"])
+
+        if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+            add_stochastic_(p.data, -update)
+        else:
+            p.data.add_(-update)
+        del update
+
+        state['step'] += 1
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for i, p in enumerate(group['params']):
+                self.step_parameter(p, group, i)
+
+        return loss

{adv_optm-0.1.6 → adv_optm-0.1.8}/adv_optm/optim/__init__.py

@@ -1,6 +1,7 @@
 from .AdamW_adv import AdamW_adv
 from .Prodigy_adv import Prodigy_adv
 from .Adopt_adv import Adopt_adv
+from .Simplified_AdEMAMix import Simplified_AdEMAMix
 from .Lion_adv import Lion_adv
 from .Lion_Prodigy_adv import Lion_Prodigy_adv
 
@@ -8,6 +9,7 @@ __all__ = [
     "AdamW_adv",
     "Prodigy_adv",
     "Adopt_adv",
+    "Simplified_AdEMAMix",
     "Lion_adv",
     "Lion_Prodigy_adv",
 ]

{adv_optm-0.1.6 → adv_optm-0.1.8}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-0.1.6 → adv_optm-0.1.8}/adv_optm.egg-info/SOURCES.txt

@@ -12,6 +12,7 @@ adv_optm/optim/Adopt_adv.py
 adv_optm/optim/Lion_Prodigy_adv.py
 adv_optm/optim/Lion_adv.py
 adv_optm/optim/Prodigy_adv.py
+adv_optm/optim/Simplified_AdEMAMix.py
 adv_optm/optim/__init__.py
 adv_optm/util/BF16_Stochastic_Rounding.py
 adv_optm/util/Effective_Shape.py

{adv_optm-0.1.6 → adv_optm-0.1.8}/setup.py

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