adv-optm 0.1.0__py3-none-any.whl

adv_optm/optim/Prodigy_adv.py

@@ -0,0 +1,367 @@
+import torch
+from typing import Optional
+import math
+
+from ..util.BF16_Stochastic_Rounding import add_stochastic_
+from ..util.Effective_Shape import _get_effective_shape
+from ..util.NNMF import _nnmf,_unnmf
+from ..util.OrthoGrad import _orthogonalize_gradient
+from ..util.One_Bit_Boolean import _pack_bools, _unpack_bools
+
+class Prodigy_adv(torch.optim.Optimizer):
+    """
+    Implements a factored Prodigy/AdamW algorithm.
+    This is an advanced version of Prodigy 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,
+        # prodigy parameters
+        beta3: float = None,
+        d0: float = 1e-6,
+        d_coef: float = 1,
+        growth_rate: float = float('inf'),
+        safeguard_warmup: bool = False,
+        slice_p: int = 11,
+    ):
+        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,
+            "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,
+        }
+        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)
+        self.init_step()
+
+    @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
+
+    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: int | None = None):
+        if p.grad is None:
+            return
+
+        grad = p.grad
+        if grad.dtype != torch.float32 and self.factored:
+            grad = grad.float()
+        if group["use_orthograd"]:
+            grad = _orthogonalize_gradient(p, grad)
+        state = self.state[p]
+
+        # State Initialization
+        if len(state) == 0:
+            state['step'] = 0
+
+            should_factor = (
+                self.factored and
+                not (len(p.shape) == 1 and not group['vector_reshape'])
+            )
+
+            state['factored'] = should_factor
+
+            slice_p = group['slice_p']
+
+            dtype = torch.float32 if self.factored else p.dtype
+            device = p.device
+
+            if state['factored']:
+                state['effective_shape'] = _get_effective_shape(p.numel())
+                d1, d2 = state['effective_shape']
+
+                # First moment (m)
+                state['mu_m_nmf'] = torch.zeros(d1, device=device, dtype=dtype) 
+                state['mv_m_nmf'] = torch.zeros(d2, device=device, dtype=dtype)
+                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
+                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)
+
+            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=device, dtype=p.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
+            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_(self.beta1).add_(grad_reshaped, alpha=self.d * (1.0 - self.beta1))
+            if self.use_grams:
+                mt.copy_(grad_reshaped.sign() * mt.abs())
+            elif self.use_cautious:
+                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_(self.beta2).addcmul_(grad_reshaped, grad_reshaped, value=self.d * self.d * (1.0 - self.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=self.d * (1.0 - beta3_ema))
+                update_m = mt + (alpha_t * mt_slow)
+            else:
+                update_m = mt
+            del grad_reshaped
+
+            if group['use_atan2']:
+                a = 1.2732395
+                denom = vt.sqrt()
+                update = torch.atan2(update_m, denom).mul_(a)
+            else:
+                denom = vt.sqrt().add_(group['eps'])
+                update = update_m / denom
+            del update_m, denom
+
+            update = update.view(p.shape)
+            update.mul_(self.dlr)
+
+            # 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 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, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+
+            exp_avg.mul_(self.beta1).add_(grad, alpha=self.d * (1.0 - self.beta1))
+            if self.use_grams:
+                exp_avg = grad.sign() * exp_avg.abs()
+            elif self.use_cautious:
+                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=self.d * (1.0 - beta3_ema))
+                update_m = exp_avg + (alpha_t * exp_avg_slow)
+            else:
+                update_m = exp_avg
+
+            exp_avg_sq.mul_(self.beta2).addcmul_(grad, grad.conj(), value=self.d * self.d * (1.0 - self.beta2))
+
+            if group['use_atan2']:
+                a = 1.2732395
+                denom = exp_avg_sq.sqrt()
+                update = torch.atan2(update_m, denom).mul_(a)
+            else:
+                denom = exp_avg_sq.sqrt().add_(group['eps'])
+                update = update_m / denom
+            del update_m, denom
+
+            update = update.mul_(self.dlr)
+
+        # --- Accumulate Prodigy stats ---
+        d0, safeguard_warmup, slice_p = group['d0'], group['safeguard_warmup'], group['slice_p']
+        s, p0 = state['s'], state['p0']
+        grad_flat = grad.flatten().float()
+        p_flat = p.data.flatten().float()
+        p0 = p0.float()
+
+        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
+
+        # 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"] * 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)
+        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)
+
+
+        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']
+        
+        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/__init__.py

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

adv_optm/util/BF16_Stochastic_Rounding.py

@@ -0,0 +1,47 @@
+import torch
+from torch import Tensor
+
+def copy_stochastic_(target: Tensor, source: Tensor):
+    """
+    Nerogar's implementation of stochastic rounding in the paper "Revisiting BFloat16 Training"
+    (https://arxiv.org/abs/2010.06192).
+    see:
+    https://github.com/pytorch/pytorch/issues/120376
+    https://github.com/Nerogar/OneTrainer/blob/daae18eaed8c0fa39289b2ff79cc2c1e08577fcb/modules/util/bf16_stochastic_rounding.py
+
+    Args:
+        target: the target tensor with dtype=bfloat16
+        source: the target tensor with dtype=float32
+    """
+    # create a random 16 bit integer
+    result = torch.randint_like(
+        source,
+        dtype=torch.int32,
+        low=0,
+        high=(1 << 16),
+    )
+
+    # add the random number to the lower 16 bit of the mantissa
+    result.add_(source.view(dtype=torch.int32))
+
+    # mask off the lower 16 bit of the mantissa
+    result.bitwise_and_(-65536)  # -65536 = FFFF0000 as a signed int32
+
+    # copy the higher 16 bit into the target tensor
+    target.copy_(result.view(dtype=torch.float32))
+
+    del result
+
+def add_stochastic_(input: Tensor, other: Tensor, alpha: float = 1.0):
+    """
+    adds other to input using stochastic rounding
+
+    Args:
+        input: the input tensor with dtype=bfloat16
+        other: the other tensor
+        alpha: a multiplier for other
+    """
+    result = other.clone() if other.dtype == torch.float32 else other.to(dtype=torch.float32)
+
+    result.add_(input, alpha=alpha)
+    copy_stochastic_(input, result)

adv_optm/util/Effective_Shape.py

@@ -0,0 +1,8 @@
+def _get_effective_shape(numel: int) -> tuple[int, int]:
+    """Finds two factors of numel that are closest to its square root."""
+    if numel <= 0:
+        return (0, 0)
+    for i in reversed(range(1, int(numel ** 0.5) + 1)):
+        if numel % i == 0:
+            return (numel // i, i)
+    return (numel, 1)

adv_optm/util/NNMF.py

@@ -0,0 +1,18 @@
+import torch
+
+def _unnmf(row_col: tuple) -> torch.Tensor:
+    """Reconstructs a matrix from its rank-1 factors (outer product)."""
+    return torch.outer(row_col[0], row_col[1])
+
+def _nnmf(matrix: torch.Tensor, out: tuple):
+    """Performs a rank-1 non-negative matrix factorization."""
+    shape = matrix.shape
+    torch.sum(matrix, dim=1, out=out[0])
+    torch.sum(matrix, dim=0, out=out[1])
+    # Normalize one of the factors for stability
+    if shape[0] < shape[1]:
+        scale = out[0].sum()
+        if scale != 0: out[0].div_(scale)
+    else:
+        scale = out[1].sum()
+        if scale != 0: out[1].div_(scale)

adv_optm/util/One_Bit_Boolean.py

@@ -0,0 +1,22 @@
+import torch
+
+@torch.no_grad()
+def _pack_bools(tensor: torch.Tensor) -> torch.Tensor:
+    """Packs a boolean tensor into a uint8 tensor to achieve 1-bit storage."""
+    n, m = tensor.shape
+    packed_m = (m + 7) // 8
+    padded_tensor = torch.nn.functional.pad(tensor, (0, packed_m * 8 - m), 'constant', 0)
+    reshaped = padded_tensor.view(n, packed_m, 8)
+    shifter = torch.arange(8, device=tensor.device, dtype=torch.uint8)
+    packed = (reshaped.to(torch.uint8) * (2**shifter)).sum(dim=2).to(torch.uint8)
+    return packed
+
+@torch.no_grad()
+def _unpack_bools(packed_tensor: torch.Tensor, original_m: int) -> torch.Tensor:
+    """Unpacks a uint8 tensor back into a boolean tensor."""
+    if packed_tensor.dtype != torch.uint8:
+        packed_tensor = packed_tensor.to(torch.uint8)
+    shifter = (2**torch.arange(8, device=packed_tensor.device, dtype=torch.uint8)).view(1, 1, 8)
+    unpacked_padded = (packed_tensor.unsqueeze(2) & shifter) != 0
+    unpacked = unpacked_padded.view(packed_tensor.shape[0], -1)[:, :original_m]
+    return unpacked

adv_optm/util/OrthoGrad.py

@@ -0,0 +1,16 @@
+import torch
+
+def _orthogonalize_gradient(p: torch.Tensor, grad: torch.Tensor) -> torch.Tensor:
+    """Projects the gradient `grad` to be orthogonal to the parameter `p`."""
+    if grad.is_sparse: raise RuntimeError("OrthoGrad logic does not support sparse gradients.")
+    original_shape = grad.shape
+    original_dtype = grad.dtype
+    w = p.view(-1).float()
+    g = grad.view(-1).float()
+    w_norm_sq = torch.dot(w, w).add_(1e-30)
+    proj = torch.dot(w, g) / w_norm_sq
+    g_orth = g.sub(w, alpha=proj)
+    g_norm = g.norm(2)
+    g_orth_norm = g_orth.norm(2).add_(1e-30)
+    g_orth_scaled = g_orth * (g_norm / g_orth_norm)
+    return g_orth_scaled.view(original_shape).to(original_dtype)

adv_optm/util/Randomized_SVD.py

@@ -0,0 +1,37 @@
+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/util/__init__.py

@@ -0,0 +1,11 @@
+from .BF16_Stochastic_Rounding import add_stochastic_, copy_stochastic_
+from .Effective_Shape import _get_effective_shape
+from .One_Bit_Boolean import _pack_bools, _unpack_bools
+from .OrthoGrad import _orthogonalize_gradient
+
+__all__ = [
+    "_pack_bools", "_unpack_bools",
+    "add_stochastic_",
+    "_get_effective_shape",
+    "_orthogonalize_gradient",
+]

adv_optm-0.1.0.dist-info/METADATA

@@ -0,0 +1,134 @@
+Metadata-Version: 2.4
+Name: adv_optm
+Version: 0.1.0
+Summary: A family of highly efficient, lightweight yet powerful optimizers.
+Home-page: https://github.com/Koratahiu/Advanced_Optimizers
+Author: Koratahiu
+Author-email: hiuhonor@gmail.com
+License: Apache 2.0
+Keywords: llm,fine-tuning,memory-efficient,low-rank,compression,pytorch,optimizer,adam
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Advanced Optimizers
+
+This repo introduces a new family of highly efficient, lightweight yet powerful optimizers, born from extensive research into recent academic literature and validated through practical training runs across diverse models.
+
+---
+
+### Install
+
+`pip install adv_optm`
+
+---
+
+### Theory (Inspired by SMMF)
+
+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).
+- 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.
+- Final storage: **four factored vectors + one 1-bit sign**.
+- Updates behave like full-state Adam but with drastically reduced memory.
+
+> ✅ **TL;DR**: Lightweight, strong, memory-efficient optimizer.
+
+---
+
+### Memory Cost
+
+- **Adopt_Factored** for full SDXL finetune: **328 MB** (4 small vectors + 1-bit state)
+- **Adopt_Factored with AdEMAMix** for full SDXL finetune: **625 MB** (6 small vectors + two 1-bit states)
+> SDXL is 6.5GB model.
+
+---
+
+### ⏱️ Speed (my tests in SDXL - BS 4)
+
+- **Adopt_Factored**: ~10s/it
+- **Adopt_Factored with AdEMAMix**: ~12s/it
+- **Adafactor**: ~8.5s/it  
+→ Overhead from compression/reconstruction cycles.
+→ It's faster than [MLorc](https://arxiv.org/abs/2506.01897) (~12s/it), which uses RSVD compression, and should be the fastest momentum compression (AFAIK).
+
+---
+
+### 📈 Performance
+
+- **Better than Adafactor, and CAME factorzation methods**
+- **Comparable or identical to Adam** (see SMMF paper results)
+
+---
+
+### Available Optimizers (all support `Factored` toggle)
+
+Set `Factored=False` to disable factorization and run as a full uncompressed optimizer (like vanilla Adam).
+
+1. **Adam**
+2. **Prodigy**
+3. **Adopt**
+
+---
+
+### Bonus Features (Built-in)
+
+- **Fused Backward Pass**
+
+- **Stochastic Rounding (SR)**: Improves quality and convergence for **BF16 training**.
+
+- **[AdEMAMix](https://arxiv.org/abs/2409.03137)**  
+  → This adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, especially during long runs of full finetuning.
+  → A higher value of beta3 (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 (2k-4k steps).
+  → When `factored` is true, it compresses the new momentum in the same way as the first moment (1-bit state + 2 vectors). However, this introduces noticeable overhead as we are compressing/reconstructing a third state each step.
+
+  ⚠️ **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.
+
+- **[OrthoGrad](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)**  
+  → Removes gradient component parallel to weights → prevents "naïve loss minimization" (NLM) → reduces natural overfitting  
+  → Perfect for fine-tuning the direction of existing features (e.g., full finetune or training a trained LoRA) without weight decay erasing prior knowledge.
+
+  ⚠️ **Note**: OrthoGrad introduces **~33% time overhead**, so take this into account.
+
+- **[Grams: Gradient Descent with Adaptive Momentum Scaling](https://github.com/Gunale0926/Grams)**  
+  → Eliminates the need for 1-bit momentum sign storage by using the **sign of gradients** for the first moment.
+
+  ⚠️ **Not recommended for small batch sizes**: gradients are too noisy, which can destabilize momentum (tested for Prodigy and it made the optimizer slower to find the LR or converge in BS 4).
+
+### Other Notes
+
+- **Adopt** skips the first step (only initializes the states) and has built-in clipping (sticking to the original optimizer), but we skip both of these when you enable `use_atan2`; as the optimizer becomes scale-invariant and the values of the states won't cause any issues or instability.
+
+- 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).
+
+---