adv-optm 0.1.3__tar.gz → 0.1.5__tar.gz

{adv_optm-0.1.3 → adv_optm-0.1.5}/PKG-INFO

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

{adv_optm-0.1.3 → adv_optm-0.1.5}/adv_optm/__init__.py

@@ -14,4 +14,4 @@ __all__ = [
     "Lion_Prodigy_adv",
 ]
 
-__version__ = "0.1.3"
+__version__ = "0.1.5"

{adv_optm-0.1.3 → adv_optm-0.1.5}/adv_optm/optim/AdamW_adv.py

@@ -1,297 +1,309 @@
-import torch
-from typing import Optional
-
-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
-
-class AdamW_adv(torch.optim.Optimizer):
-    """
-    Implements a factored AdamW algorithm.
-    This is an advanced version of AdamW with optional features like
-    low-rank factorization of optimizer states (SMMF), OrthoGrad, etc.
-
-    Args:
-        params (iterable): iterable of parameters to optimize or dicts defining
-            parameter groups
-        lr (float): learning rate (default: 1e-3)
-        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)
-        weight_decay (float): weight decay (L2 penalty) (default: 0)
-        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_atan2 (bool): whether to use the atan2 update rule. (default: False)
-        use_grams (bool): whether to use Grams-style updates. (default: False)
-        use_cautious (bool):  whether to use cautious masking to align the gradient's
-            direction with the first moment's.  (default: False)
-        use_orthograd (bool): whether to use OrthoGrad.  (default: False)
-        use_AdEMAMix (bool): whether to enable the AdEMAMix feature. This adds
-            a second, slow-moving average of the momentum (`mt_slow`) which is
-            combined with the primary momentum (`mt`) to stabilize updates,
-            especially in noisy, small-batch settings. If `False`, the
-            optimizer behaves as standard AdamW. (default: False)
-        beta3_ema (float): The decay rate for the slow exponential moving average of
-            the momentum (only used when `use_AdEMAMix` is `True`). A higher
-            value (e.g., 0.9999) gives the EMA a longer memory, making it more
-            stable but slower to adapt. A lower value (e.g., 0.999) is often
-            better for shorter training runs. (default: 0.9999)
-        alpha (float): The mixing coefficient that scales the slow momentum term
-            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)
-        t_alpha (Optional[int]): The number of steps for a linear warmup of the
-            `alpha` parameter (only used when `use_AdEMAMix` is `True`). This is
-            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
-        factored (bool): whether to use the factorization or disable it to use
-            the uncompressed optimizer. (default: True)
-    """
-
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-3,
-        betas: tuple[float, float] = (0.9, 0.999),
-        eps: float = 1e-8,
-        weight_decay: float = 0.0,
-        vector_reshape: bool = True,
-        stochastic_rounding: bool = True,
-        use_atan2: bool = False,
-        use_cautious: bool = False,
-        use_grams: bool = False,
-        use_orthograd: bool = False,
-        use_AdEMAMix: bool = False,
-        beta3_ema: float = 0.9999,
-        alpha: float = 5.0,
-        t_alpha: int | None = None,
-        factored: bool = True,
-    ):
-        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}")
-
-        defaults = {
-            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
-            "vector_reshape": vector_reshape, "use_atan2": use_atan2,
-            "use_orthograd": use_orthograd,
-            "beta3_ema": beta3_ema, "alpha": alpha, "t_alpha": t_alpha,
-        }
-        self.stochastic_rounding = stochastic_rounding
-        self.use_cautious = use_cautious
-        self.use_grams = use_grams
-        self.use_AdEMAMix = use_AdEMAMix
-        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]
-
-        beta1, beta2 = group['betas']
-
-        # 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)
-                if beta1 > 0:
-                    state['mu_m_nmf'] = torch.zeros(d1, device=device, dtype=dtype) 
-                    state['mv_m_nmf'] = torch.zeros(d2, device=device, dtype=dtype)
-                    if not self.use_grams:
-                        packed_d2 = (d2 + 7) // 8
-                        state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=device)
-                if self.use_AdEMAMix:
-                    state['mu_m_slow_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype) 
-                    state['mv_m_slow_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
-                    packed_d2 = (d2 + 7) // 8
-                    state['sign_slow'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.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 AdamW for non-factored tensors
-                if beta1 > 0:
-                    state['exp_avg'] = torch.zeros_like(p, device=device, dtype=dtype)
-                if self.use_AdEMAMix:
-                    state['exp_avg_slow'] = torch.zeros_like(p, dtype=dtype)
-                state['exp_avg_sq'] = torch.zeros_like(p, device=device, dtype=dtype)
-
-        if self.use_AdEMAMix:
-            beta3_ema = group['beta3_ema']
-            alpha = group['alpha']
-            t_alpha = group['t_alpha']
-            current_step = state['step'] + 1
-            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 state['factored']:
-            d1, d2 = state['effective_shape']
-
-            # Reconstruct momentum from previous step's factors
-            if beta1 > 0:
-                mt = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
-                if not self.use_grams:
-                    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 - beta1)
-                if self.use_grams:
-                    mt.copy_(grad_reshaped.sign() * mt.abs())
-                elif self.use_cautious:
-                    mask = (mt * grad_reshaped > 0).to(grad_reshaped.dtype)
-                    mask.div_(mask.mean().clamp_(min=1e-3))
-                    mt.mul_(mask)
-                    del mask
-
-            vt = _unnmf((state['mu_v_nmf'], state['mv_v_nmf']))
-            vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
-
-            if self.use_AdEMAMix:
-                mt_slow = _unnmf((state['mu_m_slow_nmf'], state['mv_m_slow_nmf']))
-                if state['sign_slow'].dtype != torch.uint8:
-                    state['sign_slow'] = state['sign_slow'].to(torch.uint8)
-                unpacked_sign_slow = _unpack_bools(state['sign_slow'], original_m=d2)
-                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=1.0 - beta3_ema)
-                update = mt + (alpha_t * mt_slow) if beta1 > 0 else grad_reshaped + (alpha_t * mt_slow)
-            else:
-                update = mt if beta1 > 0 else grad_reshaped
-            del grad_reshaped
-
-            if group['use_atan2']:
-                a = 1.2732395
-                denom = vt.sqrt()
-                update.atan2_(denom).mul_(a)
-            else:
-                denom = vt.sqrt()
-                update.div_(denom.add_(group['eps']))
-            del denom
-
-            update.view(p.shape).mul_(group['lr'])
-
-            # Compress updated moments and store new factors
-            if beta1 > 0:
-                if not self.use_grams:
-                    state['sign'] = _pack_bools(mt > 0)
-                _nnmf(mt.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
-                del mt
-            if self.use_AdEMAMix:
-                state['sign_slow'] = _pack_bools(mt_slow > 0)
-                _nnmf(mt_slow.abs(), out=(state['mu_m_slow_nmf'], state['mv_m_slow_nmf']))
-                del mt_slow
-            _nnmf(vt, out=(state['mu_v_nmf'], state['mv_v_nmf']))
-            del vt
-
-        else:  # Standard AdamW logic for non-factored tensors
-            exp_avg_sq = state['exp_avg_sq']
-
-            if beta1 > 0:
-                exp_avg = state['exp_avg']
-                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
-                if self.use_grams:
-                    exp_avg = grad.sign() * exp_avg.abs()
-                elif self.use_cautious:
-                    mask = (exp_avg * grad > 0).to(grad.dtype)
-                    mask.div_(mask.mean().clamp_(min=1e-3))
-                    exp_avg.mul_(mask)
-                    del mask
-
-            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)
-            else:
-                update = exp_avg if beta1 > 0 else grad
-
-            exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
-
-            if group['use_atan2']:
-                a = 1.2732395
-                denom = exp_avg_sq.sqrt()
-                update.atan2_(denom).mul_(a)
-            else:
-                denom = exp_avg_sq.sqrt()
-                update.div_(denom.add_(group['eps']))
-            del denom
-
-            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)
-
+import torch
+from typing import Optional
+
+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
+
+class AdamW_adv(torch.optim.Optimizer):
+    """
+    Implements a factored AdamW algorithm.
+    This is an advanced version of AdamW with optional features like
+    low-rank factorization of optimizer states (SMMF), OrthoGrad, etc.
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float): learning rate (default: 1e-3)
+        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)
+        weight_decay (float): weight decay (L2 penalty) (default: 0).
+        use_bias_correction (bool): whether to use bias correction for the first
+            and second moment estimates, as in the original Adam paper.
+            (default: True)
+        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_atan2 (bool): whether to use the atan2 update rule. (default: False)
+        use_grams (bool): whether to use Grams-style updates. (default: False)
+        use_cautious (bool):  whether to use cautious masking to align the gradient's
+            direction with the first moment's.  (default: False)
+        use_orthograd (bool): whether to use OrthoGrad.  (default: False)
+        use_AdEMAMix (bool): whether to enable the AdEMAMix feature. This adds
+            a second, slow-moving average of the momentum (`mt_slow`) which is
+            combined with the primary momentum (`mt`) to stabilize updates,
+            especially in noisy, small-batch settings. If `False`, the
+            optimizer behaves as standard AdamW. (default: False)
+        beta3_ema (float): The decay rate for the slow exponential moving average of
+            the momentum (only used when `use_AdEMAMix` is `True`). A higher
+            value (e.g., 0.9999) gives the EMA a longer memory, making it more
+            stable but slower to adapt. A lower value (e.g., 0.999) is often
+            better for shorter training runs. (default: 0.9999)
+        alpha (float): The mixing coefficient that scales the slow momentum term
+            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)
+        t_alpha (Optional[int]): The number of steps for a linear warmup of the
+            `alpha` parameter (only used when `use_AdEMAMix` is `True`). This is
+            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. (default: None)
+        factored (bool): whether to use the factorization or disable it to use
+            the uncompressed optimizer. (default: True)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-3,
+        betas: tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+        use_bias_correction: bool = True,
+        vector_reshape: bool = True,
+        stochastic_rounding: bool = True,
+        use_atan2: bool = False,
+        use_cautious: bool = False,
+        use_grams: bool = False,
+        use_orthograd: bool = False,
+        use_AdEMAMix: bool = False,
+        beta3_ema: float = 0.9999,
+        alpha: float = 5.0,
+        t_alpha: int | None = None,
+        factored: bool = True,
+    ):
+        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}")
+
+        defaults = {
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
+            "vector_reshape": vector_reshape, "use_atan2": use_atan2,
+            "use_orthograd": use_orthograd, "use_bias_correction": use_bias_correction,
+            "beta3_ema": beta3_ema, "alpha": alpha, "t_alpha": t_alpha,
+        }
+        self.stochastic_rounding = stochastic_rounding
+        self.use_cautious = use_cautious
+        self.use_grams = use_grams
+        self.use_AdEMAMix = use_AdEMAMix
+        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]
+
+        beta1, beta2 = group['betas']
+
+        # 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)
+                if beta1 > 0:
+                    state['mu_m_nmf'] = torch.zeros(d1, device=device, dtype=dtype) 
+                    state['mv_m_nmf'] = torch.zeros(d2, device=device, dtype=dtype)
+                    if not self.use_grams:
+                        packed_d2 = (d2 + 7) // 8
+                        state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=device)
+                if self.use_AdEMAMix:
+                    state['mu_m_slow_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype) 
+                    state['mv_m_slow_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
+                    packed_d2 = (d2 + 7) // 8
+                    state['sign_slow'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.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 AdamW for non-factored tensors
+                if beta1 > 0:
+                    state['exp_avg'] = torch.zeros_like(p, device=device, dtype=dtype)
+                if self.use_AdEMAMix:
+                    state['exp_avg_slow'] = torch.zeros_like(p, device=device, dtype=dtype)
+                state['exp_avg_sq'] = torch.zeros_like(p, device=device, dtype=dtype)
+
+        step = state['step'] + 1
+        if group['use_bias_correction']:
+            bias_correction1 = 1.0 - beta1 ** step
+            bias_correction2 = 1.0 - beta2 ** step
+        else:
+            bias_correction1 = 1
+            bias_correction2 = 1
+        step_size = group['lr'] / bias_correction1
+
+        if self.use_AdEMAMix:
+            beta3_ema = group['beta3_ema']
+            alpha = group['alpha']
+            t_alpha = group['t_alpha']
+            alpha_t = alpha
+            if t_alpha is not None and t_alpha > 0 and step < t_alpha:
+                alpha_t = min(step * alpha / t_alpha, alpha)
+
+        if state['factored']:
+            d1, d2 = state['effective_shape']
+
+            # Reconstruct momentum from previous step's factors
+            if beta1 > 0:
+                mt = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
+                if not self.use_grams:
+                    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 - beta1)
+                if self.use_grams:
+                    mt.copy_(grad_reshaped.sign() * mt.abs())
+                elif self.use_cautious:
+                    mask = (mt * grad_reshaped > 0).to(grad_reshaped.dtype)
+                    mask.div_(mask.mean().clamp_(min=1e-3))
+                    mt.mul_(mask)
+                    del mask
+
+            vt = _unnmf((state['mu_v_nmf'], state['mv_v_nmf']))
+            vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
+
+            if self.use_AdEMAMix:
+                mt_slow = _unnmf((state['mu_m_slow_nmf'], state['mv_m_slow_nmf']))
+                if state['sign_slow'].dtype != torch.uint8:
+                    state['sign_slow'] = state['sign_slow'].to(torch.uint8)
+                unpacked_sign_slow = _unpack_bools(state['sign_slow'], original_m=d2)
+                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=1.0 - beta3_ema)
+                update = mt + (alpha_t * mt_slow) if beta1 > 0 else grad_reshaped + (alpha_t * mt_slow)
+            else:
+                update = mt.clone() if beta1 > 0 else grad_reshaped.clone()
+            del grad_reshaped
+
+            if group['use_atan2']:
+                a = 1.2732395
+                denom = (vt.sqrt() / (bias_correction2**0.5))
+                update.atan2_(denom).mul_(a)
+            else:
+                denom = (vt.sqrt() / (bias_correction2**0.5)).add_(group['eps'])
+                update.div_(denom)
+            del denom
+
+            update.view(p.shape).mul_(step_size)
+
+            # Compress updated moments and store new factors
+            if beta1 > 0:
+                if not self.use_grams:
+                    state['sign'] = _pack_bools(mt > 0)
+                _nnmf(mt.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
+                del mt
+            if self.use_AdEMAMix:
+                state['sign_slow'] = _pack_bools(mt_slow > 0)
+                _nnmf(mt_slow.abs(), out=(state['mu_m_slow_nmf'], state['mv_m_slow_nmf']))
+                del mt_slow
+            _nnmf(vt, out=(state['mu_v_nmf'], state['mv_v_nmf']))
+            del vt
+
+        else:  # Standard AdamW logic for non-factored tensors
+            exp_avg_sq = state['exp_avg_sq']
+
+            if beta1 > 0:
+                exp_avg = state['exp_avg']
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+                if self.use_grams:
+                    exp_avg = grad.sign() * exp_avg.abs()
+                elif self.use_cautious:
+                    mask = (exp_avg * grad > 0).to(grad.dtype)
+                    mask.div_(mask.mean().clamp_(min=1e-3))
+                    exp_avg.mul_(mask)
+                    del mask
+
+            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)
+            else:
+                update = exp_avg.clone() if beta1 > 0 else grad.clone()
+
+            exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
+
+            if group['use_atan2']:
+                a = 1.2732395
+                denom = (exp_avg_sq.sqrt() / (bias_correction2**0.5))
+                update.atan2_(denom).mul_(a)
+            else:
+                denom = (exp_avg_sq.sqrt() / (bias_correction2**0.5)).add_(group['eps'])
+                update.div_(denom)
+            del denom
+
+            update.mul_(step_size)
+
+        # 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.3 → adv_optm-0.1.5}/adv_optm/optim/Lion_Prodigy_adv.py

@@ -216,11 +216,19 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
 
             # Update momentum m_t = β2*m_{t-1} + (1-β2)*lr*g_t
             if self.variance_reduction:
-                vr_term = grad_reshaped - state['prev_grad']
-                exp_avg.mul_(self.beta2).add_(grad_reshaped, alpha=self.d * (1-self.beta2)).add_(vr_term, alpha=self.beta2)
+                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:
-                exp_avg.mul_(self.beta2).add_(grad_reshaped, alpha=self.d * (1-self.beta2))
+                # Standard Prodigy-Lion
+                alpha = self.d * (1 - self.beta2)
+                exp_avg.mul_(self.beta2).add_(grad_reshaped, alpha=alpha)
             del grad_reshaped
 
             # Compress new momentum m_t and store factors
@@ -247,11 +255,19 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
 
             # Update momentum 
             if self.variance_reduction:
-                vr_term = grad - state['prev_grad']
-                exp_avg.mul_(self.beta2).add_(grad, alpha=self.d * (1-self.beta2)).add_(vr_term, alpha=self.beta2)
+                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:
-                exp_avg.mul_(self.beta2).add_(grad, alpha=self.d * (1-self.beta2))
+                # Standard Prodigy-Lion
+                alpha = self.d * (1 - self.beta2)
+                exp_avg.mul_(self.beta2).add_(grad, alpha=alpha)
 
         # --- Accumulate Prodigy stats ---
         d0, safeguard_warmup, slice_p = group['d0'], group['safeguard_warmup'], group['slice_p']

{adv_optm-0.1.3 → adv_optm-0.1.5}/adv_optm/optim/Lion_adv.py

@@ -1,229 +1,243 @@
-import torch
-
-from typing import Tuple, Optional
-
-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
-
-class Lion_adv(torch.optim.Optimizer):
-    """
-    Implements the SMMF technique for Lion algorithm.
-
-    This optimizer combines the Lion update rule with the memory-saving low-rank
-    compression (SMMF) technique from https://arxiv.org/abs/2412.08894.
-
-    Args:
-        params (iterable): iterable of parameters to optimize or dicts defining
-            parameter groups.
-        lr (float, optional): learning rate (default: 1e-4).
-        betas (Tuple[float, float], optional): coefficients for computing
-            running averages of the update (default: (0.9, 0.99)).
-        weight_decay (float, optional): weight decay (L2 penalty) (default: 0.0).
-        vector_reshape (bool, optional): whether to reshape 1D vectors into 2D
-            matrices to apply low-rank compression (default: True).
-        stochastic_rounding (bool, optional): whether to use stochastic
-            rounding for BF16 parameter updates (default: True).
-        use_cautious (bool): whether to use the cautious masking technique. (default: False).
-        clip_threshold (float, optional): whether to clip the gradients norm
-            per-parameter as proposed in the paper `Lions and Muons: Optimization via
-            Stochastic Frank-Wolfe` (https://arxiv.org/abs/2506.04192) to make Lion more stable
-            (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__(
-        self,
-        params,
-        lr: float = 1e-4,
-        betas: Tuple[float, float] = (0.9, 0.99),
-        weight_decay: float = 0.0,
-        vector_reshape: bool = True,
-        stochastic_rounding: bool = True,
-        use_orthograd: bool = False,
-        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}")
-        if not all(0.0 <= beta <= 1.0 for beta in betas):
-            raise ValueError(f"Betas should be in [0.0, 1.0], but got {betas}")
-        if not weight_decay >= 0.0:
-            raise ValueError(f"Weight decay must be >= 0.0, but got {weight_decay}")
-
-        defaults = dict(
-            lr=lr,
-            betas=betas,
-            weight_decay=weight_decay,
-            vector_reshape=vector_reshape,
-            use_orthograd=use_orthograd,
-            clip_threshold=clip_threshold,
-        )
-        self.stochastic_rounding = stochastic_rounding
-        self.use_cautious = use_cautious
-        self.factored = factored
-        self.variance_reduction = variance_reduction
-        super().__init__(params, defaults)
-
-    @property
-    def supports_fused_back_pass(self) -> bool:
-        return True
-
-    @property
-    def supports_memory_efficient_fp16(self) -> bool:
-        return True
-
-    @property
-    def supports_flat_params(self) -> bool:
-        return False
-
-    @torch.no_grad()
-    def step_parameter(self, p: torch.Tensor, group: dict, i: Optional[int] = None):
-        """Performs a single optimization step on a single parameter."""
-        if p.grad is None:
-            return
-
-        grad = p.grad
-        if grad.dtype != torch.float32 and self.factored:
-            grad = grad.float()
-        if group["clip_threshold"] > 0.0:
-            grad_norm = torch.norm(grad.detach())
-            if grad_norm > group["clip_threshold"]:
-                clip_coef = group["clip_threshold"] / grad_norm
-                grad.mul_(clip_coef)
-        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
-
-            if state['factored']:
-                state['effective_shape'] = _get_effective_shape(p.numel())
-                d1, d2 = state['effective_shape']
-                state['mu_m_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype) 
-                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"]
-        lr = group["lr"]
-
-        if state['factored']:
-            # Factored Path
-            d1, d2 = state['effective_shape']
-            grad_reshaped = grad.view(d1, d2)
-            # Reconstruct momentum m_{t-1}
-            exp_avg = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
-            unpacked_sign = _unpack_bools(state['sign'], original_m=d2)
-            torch.where(unpacked_sign, exp_avg, -exp_avg, out=exp_avg)
-            del unpacked_sign
-            if exp_avg.dtype != torch.float32:
-                exp_avg = exp_avg.float()
-
-            # Compute update term c_t = β1*m_{t-1} + (1-β1)*g_t
-            signed_update = exp_avg.clone().mul_(beta1).add_(grad_reshaped, alpha=(1-beta1)).sign_()
-
-            if self.use_cautious:
-                mask = (signed_update * grad_reshaped > 0).to(grad_reshaped.dtype)
-                mask.div_(mask.mean().clamp_(min=1e-3))
-                signed_update.mul_(mask)
-                del mask
-
-            # Parameter update: p_t = p_{t-1} - lr * sign(c_t)
-            update_for_param = signed_update.view(p.shape).mul_(lr)
-
-            # Update momentum m_t = β2*m_{t-1} + (1-β2)*lr*g_t
-            if self.variance_reduction:
-                vr_term = grad_reshaped - state['prev_grad']
-                exp_avg.mul_(beta2).add_(grad_reshaped, alpha=1-beta2).add_(vr_term, alpha=beta2)
-                del vr_term
-                state['prev_grad'].copy_(grad_reshaped)
-            else:
-                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)
-            _nnmf(exp_avg.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
-            del exp_avg
-
-        else:
-            # Fallback to standard Lion logic
-            exp_avg = state["exp_avg"]
-
-            # Compute update term and sign for the update
-            if exp_avg.dtype != torch.float32 and self.factored:
-                exp_avg = exp_avg.float()
-            signed_update = exp_avg.clone().mul_(beta1).add_(grad, alpha=(1-beta1)).sign_()
-
-            if self.use_cautious:
-                mask = (signed_update * grad > 0).to(grad.dtype)
-                mask.div_(mask.mean().clamp_(min=1e-3))
-                signed_update.mul_(mask)
-                del mask
-
-            update_for_param = signed_update.mul_(lr)
-
-            # Update momentum 
-            if self.variance_reduction:
-                vr_term = grad - state['prev_grad']
-                exp_avg.mul_(beta2).add_(grad, alpha=1-beta2).add_(vr_term, alpha=beta2)
-                state['prev_grad'].copy_(grad)
-            else:
-                exp_avg.mul_(beta2).add_(grad, alpha=1-beta2)
-
-        if group["weight_decay"] != 0:
-            if p.dtype == torch.bfloat16 and self.stochastic_rounding:
-                add_stochastic_(p.data, p.data,
-                                alpha=-group["weight_decay"] * lr)
-            else:
-                p.data.add_(
-                    p.data, alpha=-group["weight_decay"] * lr
-                )
-
-        if p.dtype == torch.bfloat16 and self.stochastic_rounding:
-            add_stochastic_(p.data, -update_for_param)
-        else:
-            p.data.add_(-update_for_param)
-
-            del update_for_param
-
-    @torch.no_grad()
-    def step(self, closure: Optional[callable] = 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"]):
-                if p.grad is not None:
-                    self.step_parameter(p, group, i)
-
+import torch
+
+from typing import Tuple, Optional
+
+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
+
+class Lion_adv(torch.optim.Optimizer):
+    """
+    Implements the SMMF technique for Lion algorithm.
+
+    This optimizer combines the Lion update rule with the memory-saving low-rank
+    compression (SMMF) technique from https://arxiv.org/abs/2412.08894.
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate (default: 1e-4).
+        betas (Tuple[float, float], optional): coefficients for computing
+            running averages of the update (default: (0.9, 0.99)).
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0.0).
+        vector_reshape (bool, optional): whether to reshape 1D vectors into 2D
+            matrices to apply low-rank compression (default: True).
+        stochastic_rounding (bool, optional): whether to use stochastic
+            rounding for BF16 parameter updates (default: True).
+        use_cautious (bool): whether to use the cautious masking technique. (default: False).
+        clip_threshold (float, optional): whether to clip the gradients norm
+            per-parameter as proposed in the paper `Lions and Muons: Optimization via
+            Stochastic Frank-Wolfe` (https://arxiv.org/abs/2506.04192) to make Lion more stable
+            (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__(
+        self,
+        params,
+        lr: float = 1e-4,
+        betas: Tuple[float, float] = (0.9, 0.99),
+        weight_decay: float = 0.0,
+        vector_reshape: bool = True,
+        stochastic_rounding: bool = True,
+        use_orthograd: bool = False,
+        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}")
+        if not all(0.0 <= beta <= 1.0 for beta in betas):
+            raise ValueError(f"Betas should be in [0.0, 1.0], but got {betas}")
+        if not weight_decay >= 0.0:
+            raise ValueError(f"Weight decay must be >= 0.0, but got {weight_decay}")
+
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            weight_decay=weight_decay,
+            vector_reshape=vector_reshape,
+            use_orthograd=use_orthograd,
+            clip_threshold=clip_threshold,
+        )
+        self.stochastic_rounding = stochastic_rounding
+        self.use_cautious = use_cautious
+        self.factored = factored
+        self.variance_reduction = variance_reduction
+        super().__init__(params, defaults)
+
+    @property
+    def supports_fused_back_pass(self) -> bool:
+        return True
+
+    @property
+    def supports_memory_efficient_fp16(self) -> bool:
+        return True
+
+    @property
+    def supports_flat_params(self) -> bool:
+        return False
+
+    @torch.no_grad()
+    def step_parameter(self, p: torch.Tensor, group: dict, i: Optional[int] = None):
+        """Performs a single optimization step on a single parameter."""
+        if p.grad is None:
+            return
+
+        grad = p.grad
+        if grad.dtype != torch.float32 and self.factored:
+            grad = grad.float()
+        if group["clip_threshold"] > 0.0:
+            grad_norm = torch.norm(grad.detach())
+            if grad_norm > group["clip_threshold"]:
+                clip_coef = group["clip_threshold"] / grad_norm
+                grad.mul_(clip_coef)
+        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
+
+            if state['factored']:
+                state['effective_shape'] = _get_effective_shape(p.numel())
+                d1, d2 = state['effective_shape']
+                state['mu_m_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype) 
+                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"]
+        lr = group["lr"]
+
+        if state['factored']:
+            # Factored Path
+            d1, d2 = state['effective_shape']
+            grad_reshaped = grad.view(d1, d2)
+            # Reconstruct momentum m_{t-1}
+            exp_avg = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
+            unpacked_sign = _unpack_bools(state['sign'], original_m=d2)
+            torch.where(unpacked_sign, exp_avg, -exp_avg, out=exp_avg)
+            del unpacked_sign
+            if exp_avg.dtype != torch.float32:
+                exp_avg = exp_avg.float()
+
+            # Compute update term c_t
+            signed_update = exp_avg.clone().mul_(beta1).add_(grad_reshaped, alpha=(1-beta1)).sign_()
+
+            if self.use_cautious:
+                mask = (signed_update * grad_reshaped > 0).to(grad_reshaped.dtype)
+                mask.div_(mask.mean().clamp_(min=1e-3))
+                signed_update.mul_(mask)
+                del mask
+
+            # 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)
+
+            # Compress new momentum m_t and store factors
+            state['sign'] = _pack_bools(exp_avg > 0)
+            _nnmf(exp_avg.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
+            del exp_avg
+
+        else:
+            # Fallback to standard Lion logic
+            exp_avg = state["exp_avg"]
+
+            # Compute update term and sign for the update
+            if exp_avg.dtype != torch.float32 and self.factored:
+                exp_avg = exp_avg.float()
+            signed_update = exp_avg.clone().mul_(beta1).add_(grad, alpha=(1-beta1)).sign_()
+
+            if self.use_cautious:
+                mask = (signed_update * grad > 0).to(grad.dtype)
+                mask.div_(mask.mean().clamp_(min=1e-3))
+                signed_update.mul_(mask)
+                del mask
+
+            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)
+
+        if group["weight_decay"] != 0:
+            if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+                add_stochastic_(p.data, p.data,
+                                alpha=-group["weight_decay"] * lr)
+            else:
+                p.data.add_(
+                    p.data, alpha=-group["weight_decay"] * lr
+                )
+
+        if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+            add_stochastic_(p.data, -update_for_param)
+        else:
+            p.data.add_(-update_for_param)
+
+            del update_for_param
+
+    @torch.no_grad()
+    def step(self, closure: Optional[callable] = 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"]):
+                if p.grad is not None:
+                    self.step_parameter(p, group, i)
+
         return loss

{adv_optm-0.1.3 → adv_optm-0.1.5}/adv_optm/optim/Prodigy_adv.py

@@ -265,7 +265,7 @@ class Prodigy_adv(torch.optim.Optimizer):
                 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)
             else:
-                update = mt if self.beta1 > 0 else grad_reshaped
+                update = mt.clone() if self.beta1 > 0 else grad_reshaped.clone()
             del grad_reshaped
 
             if group['use_atan2']:
@@ -311,7 +311,7 @@ class Prodigy_adv(torch.optim.Optimizer):
                 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)
             else:
-                update = exp_avg if self.beta1 > 0 else grad
+                update = exp_avg.clone() if self.beta1 > 0 else grad.clone()
 
             exp_avg_sq.mul_(self.beta2).addcmul_(grad, grad.conj(), value=self.d * self.d * (1.0 - self.beta2))
 

{adv_optm-0.1.3 → adv_optm-0.1.5}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-0.1.3 → adv_optm-0.1.5}/setup.py

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