adv-optm 1.0.0__tar.gz → 1.0.1__tar.gz

{adv_optm-1.0.0 → adv_optm-1.0.1}/PKG-INFO

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

{adv_optm-1.0.0 → adv_optm-1.0.1}/adv_optm/__init__.py

@@ -16,4 +16,4 @@ __all__ = [
     "Lion_Prodigy_adv",
 ]
 
-__version__ = "1.0.0"
+__version__ = "1.0.1"

{adv_optm-1.0.0 → adv_optm-1.0.1}/adv_optm/optim/AdamW_adv.py

@@ -30,8 +30,8 @@ class AdamW_adv(torch.optim.Optimizer):
         stochastic_rounding (bool): whether to use stochastic
             rounding for BF16 parameter updates (default: True).
         use_atan2 (bool): whether to use the atan2 update rule. (default: False)
-        use_grams (bool): whether to use Grams-style updates. (default: False)
-        use_cautious (bool):  whether to use cautious masking to align the gradient's
+        grams_moment (bool): whether to use Grams-style updates. (default: False)
+        cautious_mask (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
@@ -54,7 +54,7 @@ class AdamW_adv(torch.optim.Optimizer):
             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
+        nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
     """
 
@@ -69,14 +69,14 @@ class AdamW_adv(torch.optim.Optimizer):
         vector_reshape: bool = True,
         stochastic_rounding: bool = True,
         use_atan2: bool = False,
-        use_cautious: bool = False,
-        use_grams: bool = False,
+        cautious_mask: bool = False,
+        grams_moment: 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 = False,
+        nnmf_factor: bool = False,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -86,9 +86,9 @@ class AdamW_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 use_cautious and use_grams:
-            print("Warning: use_cautious is incompatible with use_grams, Disabling use_cautious.")
-            use_cautious = False
+        if cautious_mask and grams_moment:
+            print("Warning: cautious is incompatible with grams, Disabling cautious.")
+            cautious_mask = False
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
@@ -97,10 +97,10 @@ class AdamW_adv(torch.optim.Optimizer):
             "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.cautious_mask = cautious_mask
+        self.grams_moment = grams_moment
         self.use_AdEMAMix = use_AdEMAMix
-        self.factored = factored
+        self.factored = nnmf_factor
         super().__init__(params, defaults)
 
     @property
@@ -151,7 +151,7 @@ class AdamW_adv(torch.optim.Optimizer):
                 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:
+                    if not self.grams_moment:
                         packed_d2 = (d2 + 7) // 8
                         state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=device)
                 if self.use_AdEMAMix:
@@ -192,16 +192,16 @@ class AdamW_adv(torch.optim.Optimizer):
             # 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:
+                if not self.grams_moment:
                     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:
+                if self.grams_moment:
                     mt.copy_(grad_reshaped.sign() * mt.abs())
-                elif self.use_cautious:
+                elif self.cautious_mask:
                     mask = (mt * grad_reshaped > 0).to(grad_reshaped.dtype)
                     mask.div_(mask.mean().clamp_(min=1e-3))
                     mt.mul_(mask)
@@ -240,7 +240,7 @@ class AdamW_adv(torch.optim.Optimizer):
 
             # Compress updated moments and store new factors
             if beta1 > 0:
-                if not self.use_grams:
+                if not self.grams_moment:
                     state['sign'] = _pack_bools(mt > 0)
                 _nnmf(mt.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
                 del mt
@@ -257,9 +257,9 @@ class AdamW_adv(torch.optim.Optimizer):
             if beta1 > 0:
                 exp_avg = state['exp_avg']
                 exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
-                if self.use_grams:
+                if self.grams_moment:
                     exp_avg = grad.sign() * exp_avg.abs()
-                elif self.use_cautious:
+                elif self.cautious_mask:
                     mask = (exp_avg * grad > 0).to(grad.dtype)
                     mask.div_(mask.mean().clamp_(min=1e-3))
                     exp_avg.mul_(mask)

{adv_optm-1.0.0 → adv_optm-1.0.1}/adv_optm/optim/Adopt_adv.py

@@ -36,9 +36,9 @@ class Adopt_adv(torch.optim.Optimizer):
             rounding for BF16 parameter updates (default: True).
         use_atan2 (bool): whether to use an atan2-based normalization, which can
             improve stability by removing the need for `eps`. (default: False)
-        use_cautious (bool):  whether to use cautious masking to align the gradient's
+        cautious_mask (bool):  whether to use cautious masking to align the gradient's
             direction with the first moment's.  (default: False)
-        use_grams (bool): whether to combine the gradient's direction with the
+        grams_moment (bool): whether to combine the gradient's direction with the
             first moment's magnitude (default: False).
         use_orthograd (bool): whether to use OrthoGrad. (default: False)
         use_AdEMAMix (bool): whether to enable the AdEMAMix feature. This adds
@@ -65,14 +65,14 @@ class Adopt_adv(torch.optim.Optimizer):
         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`,
+            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)
-        factored (bool): whether to use the factorization or disable it to use
+        nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
     """
 
@@ -87,8 +87,8 @@ class Adopt_adv(torch.optim.Optimizer):
         vector_reshape: bool = True,
         stochastic_rounding: bool = True,
         use_atan2: bool = False,
-        use_cautious: bool = False,
-        use_grams: bool = False,
+        cautious_mask: bool = False,
+        grams_moment: bool = False,
         use_orthograd: bool = False,
         use_AdEMAMix: bool = False,
         beta3_ema: float = 0.9999,
@@ -96,7 +96,7 @@ class Adopt_adv(torch.optim.Optimizer):
         t_alpha: int | None = None,
         Simplified_AdEMAMix: bool = False,
         alpha_grad: float = 100.0,
-        factored: bool = False,
+        nnmf_factor: bool = False,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -106,17 +106,17 @@ class Adopt_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 use_cautious and use_grams:
-            print("Warning: use_cautious is incompatible with use_grams, Disabling use_cautious.")
-            use_cautious = False
+        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 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 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.")
         if use_atan2 and Simplified_AdEMAMix:
             print("Warning: use_atan2 is incompatible with Simplified_AdEMAMix. Disabling use_atan2.")
             use_atan2 = False
@@ -129,12 +129,12 @@ 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.use_cautious = use_cautious and not Simplified_AdEMAMix
-        self.use_grams = use_grams and not Simplified_AdEMAMix
+        self.cautious_mask = cautious_mask and not Simplified_AdEMAMix
+        self.grams_moment = grams_moment and not Simplified_AdEMAMix
         self.use_orthograd = use_orthograd
         self.use_AdEMAMix = use_AdEMAMix and not Simplified_AdEMAMix
         self.Simplified_AdEMAMix = Simplified_AdEMAMix
-        self.factored = factored
+        self.factored = nnmf_factor
         super().__init__(params, defaults)
 
     @property
@@ -176,7 +176,7 @@ class Adopt_adv(torch.optim.Optimizer):
                 # m_0 = 0
                 state['mu_m_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype) 
                 state['mv_m_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
-                if not self.use_grams:
+                if not self.grams_moment:
                     packed_d2 = (d2 + 7) // 8
                     state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
                 if self.use_AdEMAMix:
@@ -220,7 +220,7 @@ class Adopt_adv(torch.optim.Optimizer):
 
             # Reconstruct m_{t-1}
             mt = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
-            if not self.use_grams:
+            if not self.grams_moment:
                 if state['sign'].dtype != torch.uint8:
                     state['sign'] = state['sign'].to(torch.uint8)
                 unpacked_sign = _unpack_bools(state['sign'], original_m=d2)
@@ -257,9 +257,9 @@ class Adopt_adv(torch.optim.Optimizer):
                 mt.mul_(beta1).add_(normalized_grad, alpha=1.0)
             else:
                 mt.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
-            if self.use_grams:
+            if self.grams_moment:
                 mt = grad_reshaped.sign() * mt.abs()
-            elif self.use_cautious:
+            elif self.cautious_mask:
                 mask = (mt * grad_reshaped > 0).to(grad_reshaped.dtype)
                 mask.div_(mask.mean().clamp_(min=1e-3))
                 mt.mul_(mask)
@@ -284,7 +284,7 @@ class Adopt_adv(torch.optim.Optimizer):
             del grad_reshaped
 
             # Compress and store new factors
-            if not self.use_grams:
+            if not self.grams_moment:
                 state['sign'] = _pack_bools(mt > 0)
             _nnmf(mt.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
             del mt
@@ -322,9 +322,9 @@ class Adopt_adv(torch.optim.Optimizer):
             else:
                 m.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
 
-            if self.use_grams:
+            if self.grams_moment:
                 m = grad.sign() * m.abs()
-            elif self.use_cautious:
+            elif self.cautious_mask:
                 mask = (m * grad > 0).to(grad.dtype)
                 mask.div_(mask.mean().clamp_(min=1e-3))
                 m.mul_(mask)

{adv_optm-1.0.0 → adv_optm-1.0.1}/adv_optm/optim/Lion_Prodigy_adv.py

@@ -26,12 +26,12 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
             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).
+        cautious_mask (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
+        nnmf_factor (bool): whether to use the factorization or use the
             uncompressed optimizer. (default: True)
         d0 (float):
             Initial D estimate for D-adaptation (default 1e-6). Rarely needs changing.
@@ -61,9 +61,9 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
         vector_reshape: bool = True,
         stochastic_rounding: bool = True,
         use_orthograd: bool = False,
-        use_cautious: bool = False,
+        cautious_mask: bool = False,
         clip_threshold: float = 0.0,
-        factored: bool = True,
+        nnmf_factor: bool = True,
         # prodigy parameters
         beta3: float = None,
         d0: float = 1e-6,
@@ -92,8 +92,8 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
             fsdp_in_use=fsdp_in_use,
         )
         self.stochastic_rounding = stochastic_rounding
-        self.use_cautious = use_cautious
-        self.factored = factored
+        self.cautious_mask = cautious_mask
+        self.factored = nnmf_factor
         self.fsdp_in_use = fsdp_in_use
         super().__init__(params, defaults)
         # Global state for accumulating metrics across parameter updates within a single step.
@@ -197,7 +197,7 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
             # Compute update term c_t = β1*m_{t-1} + (1-β1)*g_t
             signed_update = exp_avg.clone().mul_(self.beta1).add_(grad_reshaped, alpha=(1-self.beta1)).sign_()
 
-            if self.use_cautious:
+            if self.cautious_mask:
                 mask = (signed_update * grad_reshaped > 0).to(grad_reshaped.dtype)
                 mask.div_(mask.mean().clamp_(min=1e-3))
                 signed_update.mul_(mask)
@@ -224,7 +224,7 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
                 exp_avg = exp_avg.float()
             signed_update = exp_avg.clone().mul_(self.beta1).add_(grad, alpha=(1-self.beta1)).sign_()
 
-            if self.use_cautious:
+            if self.cautious_mask:
                 mask = (signed_update * grad > 0).to(grad.dtype)
                 mask.div_(mask.mean().clamp_(min=1e-3))
                 signed_update.mul_(mask)

{adv_optm-1.0.0 → adv_optm-1.0.1}/adv_optm/optim/Lion_adv.py

@@ -26,12 +26,12 @@ class Lion_adv(torch.optim.Optimizer):
             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).
+        cautious_mask (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
+        nnmf_factor (bool): whether to use the factorization or use the
             uncompressed optimizer. (default: True)
     """
 
@@ -44,9 +44,9 @@ class Lion_adv(torch.optim.Optimizer):
         vector_reshape: bool = True,
         stochastic_rounding: bool = True,
         use_orthograd: bool = False,
-        use_cautious: bool = False,
+        cautious_mask: bool = False,
         clip_threshold: float = 0.0,
-        factored: bool = True,
+        nnmf_factor: bool = True,
     ):
         if not lr > 0.0:
             raise ValueError(f"Learning rate must be > 0.0, but got {lr}")
@@ -64,8 +64,8 @@ class Lion_adv(torch.optim.Optimizer):
             clip_threshold=clip_threshold,
         )
         self.stochastic_rounding = stochastic_rounding
-        self.use_cautious = use_cautious
-        self.factored = factored
+        self.cautious_mask = cautious_mask
+        self.factored = nnmf_factor
         super().__init__(params, defaults)
 
     @property
@@ -140,7 +140,7 @@ class Lion_adv(torch.optim.Optimizer):
             # Compute update term c_t
             signed_update = exp_avg.clone().mul_(beta1).add_(grad_reshaped, alpha=(1-beta1)).sign_()
 
-            if self.use_cautious:
+            if self.cautious_mask:
                 mask = (signed_update * grad_reshaped > 0).to(grad_reshaped.dtype)
                 mask.div_(mask.mean().clamp_(min=1e-3))
                 signed_update.mul_(mask)
@@ -167,7 +167,7 @@ class Lion_adv(torch.optim.Optimizer):
                 exp_avg = exp_avg.float()
             signed_update = exp_avg.clone().mul_(beta1).add_(grad, alpha=(1-beta1)).sign_()
 
-            if self.use_cautious:
+            if self.cautious_mask:
                 mask = (signed_update * grad > 0).to(grad.dtype)
                 mask.div_(mask.mean().clamp_(min=1e-3))
                 signed_update.mul_(mask)