adv-optm 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

adv_optm/__init__.py

@@ -2,12 +2,16 @@ from .optim import (
     AdamW_adv,
     Prodigy_adv,
     Adopt_adv,
+    Lion_adv,
+    Lion_Prodigy_adv,
 )
 
 __all__ = [
     "AdamW_adv",
     "Prodigy_adv",
     "Adopt_adv",
+    "Lion_adv",
+    "Lion_Prodigy_adv",
 ]
 
-__version__ = "0.1.0"
+__version__ = "0.1.2"

adv_optm/optim/AdamW_adv.py

@@ -22,7 +22,6 @@ class AdamW_adv(torch.optim.Optimizer):
         eps (float): term added to the denominator to improve
             numerical stability (default: 1e-8)
         weight_decay (float): weight decay (L2 penalty) (default: 0)
-        use_bias_correction (boolean): Turn on Adam's bias correction. (default: False)
         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
@@ -37,7 +36,7 @@ class AdamW_adv(torch.optim.Optimizer):
             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 (float): The decay rate for the slow exponential moving average of
+        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
@@ -63,7 +62,6 @@ class AdamW_adv(torch.optim.Optimizer):
         betas: tuple[float, float] = (0.9, 0.999),
         eps: float = 1e-8,
         weight_decay: float = 0.0,
-        use_bias_correction: bool = False,
         vector_reshape: bool = True,
         stochastic_rounding: bool = True,
         use_atan2: bool = False,
@@ -71,7 +69,7 @@ class AdamW_adv(torch.optim.Optimizer):
         use_grams: bool = False,
         use_orthograd: bool = False,
         use_AdEMAMix: bool = False,
-        beta3: float = 0.9999,
+        beta3_ema: float = 0.9999,
         alpha: float = 5.0,
         t_alpha: int | None = None,
         factored: bool = True,
@@ -88,8 +86,8 @@ class AdamW_adv(torch.optim.Optimizer):
         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": beta3, "alpha": alpha, "t_alpha": t_alpha,
+            "use_orthograd": use_orthograd,
+            "beta3_ema": beta3_ema, "alpha": alpha, "t_alpha": t_alpha,
         }
         self.stochastic_rounding = stochastic_rounding
         self.use_cautious = use_cautious
@@ -122,6 +120,8 @@ class AdamW_adv(torch.optim.Optimizer):
             grad = _orthogonalize_gradient(p, grad)
         state = self.state[p]
 
+        beta1, beta2 = group['betas']
+
         # State Initialization
         if len(state) == 0:
             state['step'] = 0
@@ -141,11 +141,12 @@ class AdamW_adv(torch.optim.Optimizer):
                 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)
-                if not self.use_grams:
-                    packed_d2 = (d2 + 7) // 8
-                    state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=device)
+                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)
@@ -155,14 +156,14 @@ class AdamW_adv(torch.optim.Optimizer):
                 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
-                state['exp_avg'] = torch.zeros_like(p, device=device, dtype=dtype)
+                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)
 
-        beta1, beta2 = group['betas']
         if self.use_AdEMAMix:
-            beta3 = group['beta3']
+            beta3_ema = group['beta3_ema']
             alpha = group['alpha']
             t_alpha = group['t_alpha']
             current_step = state['step'] + 1
@@ -174,21 +175,22 @@ class AdamW_adv(torch.optim.Optimizer):
             d1, d2 = state['effective_shape']
 
             # Reconstruct momentum from previous step's factors
-            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
+            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)
@@ -201,29 +203,29 @@ class AdamW_adv(torch.optim.Optimizer):
                 torch.where(unpacked_sign_slow, mt_slow, -mt_slow, out=mt_slow)
                 del unpacked_sign_slow
 
-                mt_slow.mul_(beta3).add_(grad_reshaped, alpha=1.0 - beta3)
-                update_m = mt + (alpha_t * mt_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_m = mt
+                update = mt if beta1 > 0 else grad_reshaped
             del grad_reshaped
 
             if group['use_atan2']:
                 a = 1.2732395
                 denom = vt.sqrt()
-                update = torch.atan2(update_m, denom).mul_(a)
+                update.atan2_(denom).mul_(a)
             else:
-                denom = vt.sqrt().add_(group['eps'])
-                update = update_m / denom
-            del update_m, denom
+                denom = vt.sqrt()
+                update.div_(denom.add_(group['eps']))
+            del denom
 
-            update = update.view(p.shape)
-            update.mul_(group['lr'])
+            update.view(p.shape).mul_(group['lr'])
 
             # Compress updated moments and store new factors
-            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 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']))
@@ -232,36 +234,38 @@ class AdamW_adv(torch.optim.Optimizer):
             del vt
 
         else:  # Standard AdamW logic for non-factored tensors
-            exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
-
-            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
+            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).add_(grad, alpha=1 - beta3)
-                update_m = exp_avg + (alpha_t * 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_m = exp_avg
+                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 = torch.atan2(update_m, denom).mul_(a)
+                update.atan2_(denom).mul_(a)
             else:
-                denom = exp_avg_sq.sqrt().add_(group['eps'])
-                update = update_m / denom
-            del update_m, denom
+                denom = exp_avg_sq.sqrt()
+                update.div_(denom.add_(group['eps']))
+            del denom
 
-            update = update.mul_(group['lr'])
+            update.mul_(group['lr'])
 
         # Decoupled weight decay
         if group["weight_decay"] != 0:

adv_optm/optim/Adopt_adv.py

@@ -30,8 +30,6 @@ class Adopt_adv(torch.optim.Optimizer):
         clip_lambda (Callable, optional): A function that takes the current step
             and returns a value to clip the normalized gradient. Only used when
             `use_atan2` is False. (default: `lambda step: step**0.25`)
-        rank (int): the rank for the low-rank approximation (default: 4).
-        oversampling (int): oversampling parameter for Randomized SVD. (default: 0).
         vector_reshape (bool): whether to reshape 1D vectors into 2D
             matrices for low-rank compression (default: True).
         stochastic_rounding (bool): whether to use stochastic
@@ -48,7 +46,7 @@ class Adopt_adv(torch.optim.Optimizer):
             combined with the primary momentum (`mt`) to stabilize updates,
             especially in noisy, small-batch settings. If `False`, the
             optimizer behaves as standard ADOPT. (default: False)
-        beta3 (float): The decay rate for the slow exponential moving average of
+        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
@@ -83,7 +81,7 @@ class Adopt_adv(torch.optim.Optimizer):
         use_grams: bool = False,
         use_orthograd: bool = False,
         use_AdEMAMix: bool = False,
-        beta3: float = 0.9999,
+        beta3_ema: float = 0.9999,
         alpha: float = 5.0,
         t_alpha: int | None = None,
         factored: bool = True,
@@ -99,7 +97,7 @@ class Adopt_adv(torch.optim.Optimizer):
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
-            "vector_reshape": vector_reshape, "beta3": beta3, "alpha": alpha,
+            "vector_reshape": vector_reshape, "beta3_ema": beta3_ema, "alpha": alpha,
             "t_alpha": t_alpha,
         }
         self.clip_lambda = clip_lambda
@@ -179,7 +177,7 @@ class Adopt_adv(torch.optim.Optimizer):
 
         beta1, beta2 = group['betas']
         if self.use_AdEMAMix:
-            beta3 = group['beta3']
+            beta3_ema = group['beta3_ema']
             alpha = group['alpha']
             t_alpha = group['t_alpha']
             # Use step+1 for 1-based step count in scheduler
@@ -192,29 +190,29 @@ class Adopt_adv(torch.optim.Optimizer):
             d1, d2 = state['effective_shape']
 
             # Reconstruct m_{t-1}
-            mt_prev = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
+            mt = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
             if not self.use_grams:
                 if state['sign'].dtype != torch.uint8:
                     state['sign'] = state['sign'].to(torch.uint8)
                 unpacked_sign = _unpack_bools(state['sign'], original_m=d2)
-                torch.where(unpacked_sign, mt_prev, -mt_prev, out=mt_prev)
+                torch.where(unpacked_sign, mt, -mt, out=mt)
                 del unpacked_sign
 
             # Reconstruct AdEMAMix EMA
             if self.use_AdEMAMix:
-                mt_slow_prev = _unnmf((state['mu_m_slow_nmf'], state['mv_m_slow_nmf']))
+                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_prev, -mt_slow_prev, out=mt_slow_prev)
+                torch.where(unpacked_sign_slow, mt_slow, -mt_slow, out=mt_slow)
                 del unpacked_sign_slow
 
             # Reconstruct v_{t-1} using NNMF
-            vt_prev = _unnmf((state['mu_v_nmf'], state['mv_v_nmf']))
+            vt = _unnmf((state['mu_v_nmf'], state['mv_v_nmf']))
 
             # ADOPT Step A: Decorrelate g_t using v_{t-1}
             grad_reshaped = grad.view(d1, d2)
-            denom = vt_prev.sqrt()
+            denom = vt.sqrt()
 
             if self.use_atan2:
                 normalized_grad = torch.atan2(grad_reshaped, denom)
@@ -226,7 +224,7 @@ class Adopt_adv(torch.optim.Optimizer):
             del denom
 
             # ADOPT Step B: Update momentum m_t using normalized gradient
-            mt = mt_prev.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
+            mt.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
             if self.use_grams:
                 mt = grad_reshaped.sign() * mt.abs()
             elif self.use_cautious:
@@ -236,7 +234,7 @@ class Adopt_adv(torch.optim.Optimizer):
                 del mask
 
             if self.use_AdEMAMix:
-                mt_slow = mt_slow_prev.mul_(beta3).add_(normalized_grad, alpha=1.0 - beta3)
+                mt_slow.mul_(beta3_ema).add_(normalized_grad, alpha=1.0 - beta3_ema)
                 update = mt + (alpha_t * mt_slow)
                 update = update.view(p.shape)
             else:
@@ -248,20 +246,23 @@ class Adopt_adv(torch.optim.Optimizer):
                 update.mul_(group['lr']) 
 
             # Update second moment v_t for the *next* step using raw g_t
-            vt_updated = vt_prev.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
+            vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
             del grad_reshaped
 
             # Compress and store new factors
             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
 
             # factorize v_t using NMF compression
-            _nnmf(vt_updated, out=(state['mu_v_nmf'], state['mv_v_nmf']))
+            _nnmf(vt, out=(state['mu_v_nmf'], state['mv_v_nmf']))
+            del vt
 
         else: # Standard ADOPT logic for non-factored tensors
             m, v = state['exp_avg'], state['exp_avg_sq'] # m_{t-1}, v_{t-1}
@@ -293,7 +294,7 @@ class Adopt_adv(torch.optim.Optimizer):
                 del mask
 
             if self.use_AdEMAMix:
-                m_slow.mul_(beta3).add_(normalized_grad, alpha=1.0 - beta3)
+                m_slow.mul_(beta3_ema).add_(normalized_grad, alpha=1.0 - beta3_ema)
                 update = m + (alpha_t * m_slow)
             else:
                 update = m

adv_optm/optim/Lion_Prodigy_adv.py

@@ -0,0 +1,335 @@
+import torch
+import torch.distributed as dist
+
+import math
+
+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_Prodigy_adv(torch.optim.Optimizer):
+    """
+    Implements the SMMF technique and Prodigy D-Adaptation method for Lion algorithm.
+
+    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).
+        d0 (float):
+            Initial D estimate for D-adaptation (default 1e-6). Rarely needs changing.
+        d_coef (float):
+            Coefficient in the expression for the estimate of d (default 1.0).
+            Values such as 0.5 and 2.0 typically work as well. 
+            Changing this parameter is the preferred way to tune the method.
+        growth_rate (float):
+            prevent the D estimate from growing faster than this multiplicative rate.
+            Default is inf, for unrestricted. Values like 1.02 give a kind of learning
+            rate warmup effect.
+        fsdp_in_use (bool):
+            If you're using sharded parameters, this should be set to True. The optimizer
+            will attempt to auto-detect this, but if you're using an implementation other
+            than PyTorch's builtin version, the auto-detection won't work.
+        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).
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1,
+        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,
+        # prodigy parameters
+        beta3: float = None,
+        d0: float = 1e-6,
+        d_coef: float = 1,
+        growth_rate: float = float('inf'),
+        safeguard_warmup: bool = False,
+        fsdp_in_use: bool = False,
+        slice_p: int = 11,
+    ):
+        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}")
+        if variance_reduction and use_cautious:
+            print("Warning: Using both 'variance_reduction' and 'use_cautious' is not recommended and may lead to unintended effects.")
+
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            weight_decay=weight_decay,
+            vector_reshape=vector_reshape,
+            use_orthograd=use_orthograd,
+            clip_threshold=clip_threshold,
+            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,
+        )
+        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.
+        self.init_step()
+
+    @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
+
+    def init_step(self):
+        """Resets accumulators and calculates dlr for the upcoming step."""
+        self.d_denom = 0.0
+        
+        g_group = self.param_groups[0]
+        self.beta1, self.beta2 = g_group['betas']
+        self.beta3 = g_group['beta3']
+        if self.beta3 is None:
+            self.beta3 = math.sqrt(self.beta2)
+        
+        k = g_group['k']
+        self.d = g_group['d']
+        lr = g_group['lr']
+
+        self.dlr = self.d * lr
+
+        self.d_numerator = g_group.get('d_numerator', 0.0) * self.beta3
+
+    @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
+
+        if hasattr(p, "_fsdp_flattened"):
+            self.fsdp_in_use = True
+
+        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
+
+            slice_p = group['slice_p']
+
+            # D-Adaptation states
+            state['s'] = torch.zeros_like(p.flatten()[::slice_p]).detach()
+            if p.any():
+                state['p0'] = p.flatten()[::slice_p].detach().clone()
+            else:
+                state['p0'] = torch.tensor(0, device=p.device, dtype=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)
+
+        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_(self.beta1).add_(grad_reshaped, alpha=(1-self.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(self.dlr)
+
+            # 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)
+                state['prev_grad'].copy_(grad_reshaped)
+            else:
+                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
+            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_(self.beta1).add_(grad, alpha=(1-self.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(self.dlr)
+
+            # 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)
+                state['prev_grad'].copy_(grad)
+            else:
+                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']
+        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()
+
+        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
+
+        if group["weight_decay"] != 0:
+            if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+                add_stochastic_(p.data, p.data,
+                                alpha=-group["weight_decay"] * self.dlr)
+            else:
+                p.data.add_(
+                    p.data, alpha=-group["weight_decay"] * self.dlr
+                )
+
+        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)
+
+
+        self.calculate_d()
+        self.init_step()
+        return loss
+
+    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)
+
+        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/optim/Lion_adv.py

@@ -0,0 +1,231 @@
+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}")
+        if variance_reduction and use_cautious:
+            print("Warning: Using both 'variance_reduction' and 'use_cautious' is not recommended and may lead to unintended effects.")
+
+        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)
+
+        return loss

adv_optm/optim/Prodigy_adv.py

@@ -1,5 +1,6 @@
 import torch
-from typing import Optional
+import torch.distributed as dist
+
 import math
 
 from ..util.BF16_Stochastic_Rounding import add_stochastic_
@@ -17,7 +18,7 @@ class Prodigy_adv(torch.optim.Optimizer):
     Args:
         params (iterable): iterable of parameters to optimize or dicts defining
             parameter groups
-        lr (float): learning rate (default: 1e-3)
+        lr (float): learning rate (default: 1)
         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
@@ -54,12 +55,29 @@ class Prodigy_adv(torch.optim.Optimizer):
             the scheduler is disabled and th
         factored (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: True)
+        d0 (float):
+            Initial D estimate for D-adaptation (default 1e-6). Rarely needs changing.
+        d_coef (float):
+            Coefficient in the expression for the estimate of d (default 1.0).
+            Values such as 0.5 and 2.0 typically work as well. 
+            Changing this parameter is the preferred way to tune the method.
+        growth_rate (float):
+            prevent the D estimate from growing faster than this multiplicative rate.
+            Default is inf, for unrestricted. Values like 1.02 give a kind of learning
+            rate warmup effect.
+        fsdp_in_use (bool):
+            If you're using sharded parameters, this should be set to True. The optimizer
+            will attempt to auto-detect this, but if you're using an implementation other
+            than PyTorch's builtin version, the auto-detection won't work.
+        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).
     """
 
     def __init__(
         self,
         params,
-        lr: float = 1e-3,
+        lr: float = 1,
         betas: tuple[float, float] = (0.9, 0.999),
         eps: float = 1e-8,
         weight_decay: float = 0.0,
@@ -80,6 +98,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         d_coef: float = 1,
         growth_rate: float = float('inf'),
         safeguard_warmup: bool = False,
+        fsdp_in_use: bool = False,
         slice_p: int = 11,
     ):
         if not (lr >= 0.0):
@@ -98,12 +117,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,
         }
         self.stochastic_rounding = stochastic_rounding
         self.use_cautious = use_cautious
         self.use_grams = use_grams
         self.use_AdEMAMix = use_AdEMAMix
         self.factored = factored
+        self.fsdp_in_use = fsdp_in_use
         super().__init__(params, defaults)
         self.init_step()
 
@@ -142,6 +163,9 @@ class Prodigy_adv(torch.optim.Optimizer):
         if p.grad is None:
             return
 
+        if hasattr(p, "_fsdp_flattened"):
+            self.fsdp_in_use = True
+
         grad = p.grad
         if grad.dtype != torch.float32 and self.factored:
             grad = grad.float()
@@ -246,7 +270,7 @@ class Prodigy_adv(torch.optim.Optimizer):
                 denom = vt.sqrt()
                 update = torch.atan2(update_m, denom).mul_(a)
             else:
-                denom = vt.sqrt().add_(group['eps'])
+                denom = vt.sqrt().add_(self.d * group['eps'])
                 update = update_m / denom
             del update_m, denom
 
@@ -291,7 +315,7 @@ class Prodigy_adv(torch.optim.Optimizer):
                 denom = exp_avg_sq.sqrt()
                 update = torch.atan2(update_m, denom).mul_(a)
             else:
-                denom = exp_avg_sq.sqrt().add_(group['eps'])
+                denom = exp_avg_sq.sqrt().add_(self.d * group['eps'])
                 update = update_m / denom
             del update_m, denom
 
@@ -349,8 +373,16 @@ class Prodigy_adv(torch.optim.Optimizer):
         g_group = self.param_groups[0]
         d_max, d_coef, growth_rate = g_group['d_max'], g_group['d_coef'], g_group['growth_rate']
         
-        global_d_numerator = self.d_numerator
-        global_d_denom = self.d_denom
+        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:

adv_optm/optim/__init__.py

@@ -1,9 +1,13 @@
 from .AdamW_adv import AdamW_adv
 from .Prodigy_adv import Prodigy_adv
 from .Adopt_adv import Adopt_adv
+from .Lion_adv import Lion_adv
+from .Lion_Prodigy_adv import Lion_Prodigy_adv
 
 __all__ = [
     "AdamW_adv",
     "Prodigy_adv",
     "Adopt_adv",
+    "Lion_adv",
+    "Lion_Prodigy_adv",
 ]

adv_optm/util/__init__.py

@@ -1,4 +1,4 @@
-from .BF16_Stochastic_Rounding import add_stochastic_, copy_stochastic_
+from .BF16_Stochastic_Rounding import add_stochastic_
 from .Effective_Shape import _get_effective_shape
 from .One_Bit_Boolean import _pack_bools, _unpack_bools
 from .OrthoGrad import _orthogonalize_gradient

{adv_optm-0.1.0.dist-info → adv_optm-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 0.1.0
+Version: 0.1.2
 Summary: A family of highly efficient, lightweight yet powerful optimizers.
 Home-page: https://github.com/Koratahiu/Advanced_Optimizers
 Author: Koratahiu
@@ -47,7 +47,7 @@ Based primarily on:
 **[SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization](https://arxiv.org/abs/2412.08894)**
 
 The core innovation:
-- Uses fast, non-negative matrix factorization (rank 1, à la Adafactor), but **reconstructs the full state before each update** to preserve momentum accuracy, then re-factors afterward (factor → reconstruct → update → factor cycle).
+- Uses fast, non-negative matrix factorization (NNMF - rank 1), but **reconstructs the full state before each update** to preserve momentum accuracy, then re-factors afterward (factor → reconstruct → update → factor cycle).
 - For the *signed first moment*, we split into **sign + absolute value**:
   - Sign is stored as **1-bit state** via bitwise ops (SMMF originally used 8-bit with 7 bits wasted).
   - Absolute value goes through the factor/reconstruct cycle using two factored vectors + the signed state.
@@ -106,8 +106,6 @@ Set `Factored=False` to disable factorization and run as a full uncompressed opt
 
   ⚠️ **Note**: AdEMAMix updates are more aggressive than normal Adam/Adopt, so use a x2-x5 smaller LR than usual (or use Prodigy).
 
-  ⚠️ **Note**: The factored AdEMAMix is **Experimental** (as it needs more tests and validation, but it should work). Also, Adopt with AdEMAMix is **Experimental** (as Adopt normalizes the gradients for the momentum).
-
 - **[`atan2` smoothing & scaling](https://github.com/lucidrains/adam-atan2-pytorch)**  
   → Robust `eps` replacement (no tuning!) + built-in gradient clipping  
   → *Ideal for ADOPT* (which normally needs higher `eps` and clipping), so `use_atan2` is all-in-one for it.
@@ -129,6 +127,4 @@ Set `Factored=False` to disable factorization and run as a full uncompressed opt
 
 - When `use_atan2` is True, `eps` will be ignored and you should also disable any gradient clipping.
 
-- I don't recommend using **OrthoGrad** for training LoRA or embeddings, as their weights are zero-initialized and using weight decay for them should be safe and also beneficial (OrthoGrad is intended for fine-tuning pretrained models with no weight decay).
-
 ---

adv_optm-0.1.2.dist-info/RECORD

@@ -0,0 +1,18 @@
+adv_optm/__init__.py,sha256=BNYlxkuU8MFsWSY1_PLzp2XBSzpt-sxhnVuWVKRZGZ8,252
+adv_optm/optim/AdamW_adv.py,sha256=_4Vt79EB18rnIkHttA0CdMpli8sZ5f03pesdrwT5K58,12887
+adv_optm/optim/Adopt_adv.py,sha256=rzBWfFOPrMuC6vwETsw7QPKmVXcv4IJRDCTj-6eU1Qk,14798
+adv_optm/optim/Lion_Prodigy_adv.py,sha256=ql6506h_IIZvTPdGYrQdd6iEhCXHTMntqmg739fc_dw,14102
+adv_optm/optim/Lion_adv.py,sha256=jOoRbJ6u9HCK7IBI9ILOCcwprKIGTUNvUzhRd99WJK0,9410
+adv_optm/optim/Prodigy_adv.py,sha256=InR50MoE32zG6qgEkg_JzXl7uXAVRy4EYG0JDl4eKok,17324
+adv_optm/optim/__init__.py,sha256=e5UighM92LDvDB2JJwj8gDsTpXEedpytScwqS6F2FR8,300
+adv_optm/util/BF16_Stochastic_Rounding.py,sha256=Q5H0BcogmE4atP65dLoI21HKSf50lRdsBDfeF6v9Tbg,1548
+adv_optm/util/Effective_Shape.py,sha256=TBvIk1V8IuTbbBsxuekJA4e_v8JlR5Nujtut8RTWAm4,318
+adv_optm/util/NNMF.py,sha256=yRf5IP5Sjq0Uf0DxN0Q8NxEGSdD-f1ULziLVDOjY8K4,639
+adv_optm/util/One_Bit_Boolean.py,sha256=Wat49esdwohuN-OHOFMW8D0aOQgV9cP5Rl8z6yfmpos,1068
+adv_optm/util/OrthoGrad.py,sha256=NzInuBQGy_Ja__M1R9XbvqVaQ0fhGbtGgFE9YON7B3I,707
+adv_optm/util/__init__.py,sha256=qoyIF0jcLjs_vSEcsv36clw5LFNBEbifyXrrVxMH-G4,349
+adv_optm-0.1.2.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-0.1.2.dist-info/METADATA,sha256=iV5GBWtl4WphBeSIIsUoq1ay6-GJGnDD3XF6aSWWrqg,5846
+adv_optm-0.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+adv_optm-0.1.2.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-0.1.2.dist-info/RECORD,,

adv_optm/util/Randomized_SVD.py

@@ -1,37 +0,0 @@
-import torch
-from torch import Tensor
-
-from typing import Tuple
-
-def _rsvd(A: torch.Tensor, rank: int, oversampling: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Performs Randomized SVD."""
-    orig_dtype, device, (m, n) = A.dtype, A.device, A.shape
-    A_float = A.float()
-    l, true_rank = rank + oversampling, min(m, n, rank)
-    
-    if true_rank == 0:
-        return (
-            torch.zeros(m, rank, dtype=orig_dtype, device=device),
-            torch.zeros(rank, dtype=orig_dtype, device=device),
-            torch.zeros(rank, n, dtype=orig_dtype, device=device),
-        )
-    
-    if l >= min(m, n):  # Fallback to full SVD
-        U_full, S_full, Vh_full = torch.linalg.svd(A_float, full_matrices=False)
-        U, S, Vh = U_full[:, :true_rank], S_full[:true_rank], Vh_full[:true_rank, :]
-    else:  # Standard RSVD path
-        Omega = torch.randn(n, l, dtype=A_float.dtype, device=device)
-        Y = A_float @ Omega
-        Q, _ = torch.linalg.qr(Y.float())
-        B = Q.T @ A_float
-        U_tilde, S, Vh = torch.linalg.svd(B.float(), full_matrices=False)
-        U, S, Vh = (Q @ U_tilde)[:, :true_rank], S[:true_rank], Vh[:true_rank, :]
-        
-    if true_rank < rank: # Pad factors with zeros
-        U_padded = torch.zeros(m, rank, dtype=A_float.dtype, device=device)
-        S_padded = torch.zeros(rank, dtype=A_float.dtype, device=device)
-        Vh_padded = torch.zeros(rank, n, dtype=A_float.dtype, device=device)
-        U_padded[:, :true_rank], S_padded[:true_rank], Vh_padded[:true_rank, :] = U, S, Vh
-        U, S, Vh = U_padded, S_padded, Vh_padded
-        
-    return U.to(orig_dtype), S.to(orig_dtype), Vh.to(orig_dtype)

adv_optm-0.1.0.dist-info/RECORD

@@ -1,17 +0,0 @@
-adv_optm/__init__.py,sha256=4JNXqWmFkMvsUIQorZLy43BbyqZiJxMRQkCCr09sPKw,172
-adv_optm/optim/AdamW_adv.py,sha256=cvCl3bRfkENfbXwfdzZZ8k3AJ_tNx-c5kBgaguf5fnQ,12689
-adv_optm/optim/Adopt_adv.py,sha256=jDmz2Fky2t5Gv9VY5UzltF5b5TDtY3xS5pNlnj-Eox4,14952
-adv_optm/optim/Prodigy_adv.py,sha256=FFATlt4VFb7o3UocP_W4KjBIzJa_0ncsji7BsFFU_9E,15482
-adv_optm/optim/__init__.py,sha256=kX9MQhLQZGlKFPCGLXsZtooigs4wXULTEmNSSOJvcCY,178
-adv_optm/util/BF16_Stochastic_Rounding.py,sha256=Q5H0BcogmE4atP65dLoI21HKSf50lRdsBDfeF6v9Tbg,1548
-adv_optm/util/Effective_Shape.py,sha256=TBvIk1V8IuTbbBsxuekJA4e_v8JlR5Nujtut8RTWAm4,318
-adv_optm/util/NNMF.py,sha256=yRf5IP5Sjq0Uf0DxN0Q8NxEGSdD-f1ULziLVDOjY8K4,639
-adv_optm/util/One_Bit_Boolean.py,sha256=Wat49esdwohuN-OHOFMW8D0aOQgV9cP5Rl8z6yfmpos,1068
-adv_optm/util/OrthoGrad.py,sha256=NzInuBQGy_Ja__M1R9XbvqVaQ0fhGbtGgFE9YON7B3I,707
-adv_optm/util/Randomized_SVD.py,sha256=TFG417hh1t5f1n_mChnbgdQhpMoi37O04xVCe8wz8Qc,1708
-adv_optm/util/__init__.py,sha256=3yYKo23JDfHDZdGcjrDKxH8nYjk5KDB-i44kW-J4sPk,367
-adv_optm-0.1.0.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-0.1.0.dist-info/METADATA,sha256=ig2YmYzdS6DmX0KEIGkdwX-n9eciG2S2aZYog1feqmE,6342
-adv_optm-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-0.1.0.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-0.1.0.dist-info/RECORD,,