adv-optm 2.3.dev1__tar.gz → 2.3.dev3__tar.gz

{adv_optm-2.3.dev1 → adv_optm-2.3.dev3}/PKG-INFO

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

{adv_optm-2.3.dev1 → adv_optm-2.3.dev3}/adv_optm/__init__.py

@@ -28,4 +28,4 @@ __all__ = [
     "Stiefel_LoRA",
 ]
 
-__version__ = "2.3.dev1"
+__version__ = "2.3.dev3"

adv_optm-2.3.dev3/adv_optm/stiefel_optm/Stiefel_LoRA.py

@@ -0,0 +1,231 @@
+import torch
+
+from typing import Optional
+
+from ..util import param_update
+from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state
+from . import stiefel_util
+
+
+class Stiefel_LoRA(torch.optim.Optimizer):
+    """
+    Implements an advanced Stiefel_LoRA algorithm.
+
+    based on the paper:
+    ""
+    In disguise it's modified SignSGD with momentum (SignUM).
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate (default: 1e-4).
+        momentum (float, optional): coefficients for computing
+            running average of the gradients (default: 0.9).
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0.0).
+        cautious_wd (bool): Enables Cautious Weight Decay. If True, weight decay is
+            applied only to parameter coordinates where the sign of the parameter
+            and the sign of the optimizer update align (default: False).
+        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).
+        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. (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)
+        nnmf_factor (bool): whether to use the factorization or use the
+            uncompressed optimizer. (default: True)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-4,
+        momentum: float = 0.9,
+        # Decoupled/cautious weight decay
+        weight_decay: float = 0.0,
+        cautious_wd: bool = False,
+        # Stochastic Rounding for BF16
+        stochastic_rounding: bool = True,
+        # Simplified_AdEMAMix
+        alpha_grad: float = 1.0,
+        Simplified_AdEMAMix: bool = False,
+        # SMMF factorization
+        nnmf_factor: bool = False,
+        vector_reshape: bool = False,
+        # torch.compile
+        compiled_optimizer: bool = False,
+    ):
+        if not lr > 0.0:
+            raise ValueError(f"Learning rate must be > 0.0, but got {lr}")
+        if not 0.0 <= momentum <= 1.0:
+            raise ValueError(f"momentum should be in [0.0, 1.0], but got {momentum}")
+        if not weight_decay >= 0.0:
+            raise ValueError(f"Weight decay must be >= 0.0, but got {weight_decay}")
+
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            weight_decay=weight_decay,
+            cautious_wd=cautious_wd,
+            vector_reshape=vector_reshape,
+            alpha_grad=alpha_grad,
+            Simplified_AdEMAMix=Simplified_AdEMAMix,
+            nnmf_factor=nnmf_factor,
+        )
+        self.stochastic_rounding = stochastic_rounding
+        self._init_lr = lr
+
+        super().__init__(params, defaults)
+
+        if self.stochastic_rounding:
+            # For deterministic stochastic rounding, we need to seed the generator
+            # for each device used by the parameters.
+            devices = {p.device for group in self.param_groups for p in group['params'] if p.dtype == torch.bfloat16}
+            for device in devices:
+                param_update.set_seed(device)
+
+        # Initialize compiled function
+        self._compiled_step_parameter = None
+        if compiled_optimizer:
+            self.compile(fullgraph=True)
+
+    @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
+        state = self.state[p]
+
+        # State Initialization
+        if group["momentum"] > 0 and len(state) == 0:
+            state['factored'] = (
+                group['nnmf_factor'] and
+                not (len(p.shape) == 1 and not group['vector_reshape'])
+            )
+
+            dtype = torch.float32 if state['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)
+            else:
+                state['exp_avg'] = torch.zeros_like(p, device=p.device, dtype=dtype)
+
+        lr = group["lr"]
+
+        random_int_tensor = None
+
+        if group.get('compiled_optimizer', False):
+            if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+                # Pre-generate random tensor for stochastic rounding if needed.
+                random_int_tensor = param_update._get_random_int_for_sr(p)
+            lr = torch.as_tensor(lr, dtype=torch.float64)
+            step_param_fn = self._compiled_step_parameter
+        else:
+            step_param_fn = self._step_parameter
+
+        step_param_fn(p, grad, state, group, lr, random_int_tensor)
+
+    def _step_parameter(self, p, grad, state, group, lr, random_int_tensor):
+        if grad.dtype != torch.float32 and state['factored']:
+            grad = grad.float()
+
+        is_stiefel, is_stiefel_euclidean, is_scale = stiefel_util.set_flags_AB(p)
+
+        momentum = group["momentum"]
+        Simplified_AdEMAMix = group["Simplified_AdEMAMix"]
+        alpha_grad = group["alpha_grad"]
+
+        if state.get('factored'):
+            # Factored Path
+            d1, d2 = state['effective_shape']
+            grad_reshaped = grad.view(d1, d2)
+
+            if momentum > 0:
+                # Reconstruct momentum m_{t-1}
+                exp_avg = _reconstruct_state((state['mu_m_nmf'], state['mv_m_nmf'], state['sign'], d2), signed=True)
+                exp_avg.mul_(momentum).add_(grad_reshaped)
+
+                if Simplified_AdEMAMix:
+                    raw_update = exp_avg + (grad_reshaped * alpha_grad)
+                else:
+                    raw_update = exp_avg.clone()
+
+                # Compress new momentum m_t and store factors
+                state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(exp_avg, signed=True)
+            else:
+                raw_update = grad_reshaped.clone()
+
+            raw_update = raw_update.view(p.shape)
+
+        else:
+            # Fallback to standard SignSGD logic
+            if momentum > 0:
+                exp_avg = state["exp_avg"]
+                exp_avg.mul_(momentum).add_(grad)
+
+                if Simplified_AdEMAMix:
+                    raw_update = exp_avg + (grad * alpha_grad)
+                else:
+                    raw_update = exp_avg.clone()
+            else:
+                raw_update = grad.clone()
+
+        if is_stiefel:
+            update = raw_update
+        else:
+            update = raw_update.sign_()
+
+        if not is_stiefel:
+            # The RMS of a sign() tensor is 1.
+            # We scale it to 0.2 to match the approx RMS of AdamW
+            update = update.mul_(lr * 0.2)
+
+        stiefel_util.apply_stiefel_update(
+            self, p, group, update, lr,
+            random_int_tensor=random_int_tensor, 
+            is_B=is_stiefel,
+            is_A=is_stiefel_euclidean,
+            is_scale=is_scale
+        )
+
+    def compile(self, *args, **kwargs):
+        self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)
+
+    @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-2.3.dev3/adv_optm/stiefel_optm/__init__.py

@@ -0,0 +1,5 @@
+from .Stiefel_LoRA import Stiefel_LoRA
+
+__all__ = [
+    "Stiefel_LoRA",
+]

adv_optm-2.3.dev3/adv_optm/stiefel_optm/stiefel_util.py

@@ -0,0 +1,149 @@
+import torch
+
+from torch import Tensor
+
+from typing import Dict, Any
+
+def trangent_proj(p, update, lr):
+    """
+    [Stiefel-LoRA] Step 1: Tangent Space Projection
+    Formula: update = update - p @ sym(p.T @ update)
+    """
+    pt_u = torch.matmul(p.t(), update)
+    sym_pt_u = 0.5 * (pt_u + pt_u.t())
+    # Project the update onto the tangent space
+    update.sub_(torch.matmul(p, sym_pt_u))
+    del pt_u, sym_pt_u
+    rms_rescaling(update, lr)
+    return update
+
+def qr_retraction(p):
+    """[Stiefel-LoRA] Step 2: Manifold Retraction (QR Decomposition)"""
+    Q, R = torch.linalg.qr(p)
+    d = R.diagonal().sign_()
+    Q *= d
+    return p.copy_(Q)
+
+def rms_rescaling(update, lr):
+    """Rescales update to have RMS (Root Mean Square) of 0.2."""
+    # RMS = sqrt(mean(update**2))
+    # Frobenius Norm = sqrt(sum(update**2))
+    # Relationship: Norm = RMS * sqrt(numel)
+    numel = update.numel()
+    norm = torch.linalg.vector_norm(update).clamp_min_(1e-12)
+    target_norm = lr * (numel ** 0.5) * 0.2
+    return update.mul_(target_norm / norm)
+
+def set_flags_AB(p):
+    """
+    Identify if parameter is LoRA A, B, or a scale parameter.
+    """
+    if getattr(p, '_is_dora_scale', False):
+        return False, False, True
+    if getattr(p, '_is_lora_B', False):
+        return True, False, False
+    if getattr(p, '_is_lora_A', False):
+        return False, True, False
+
+    # Fallback heuristic (handles 4D Conv2d layers properly)
+    dim0 = p.shape[0]
+    dim1 = p.shape[1] if p.ndim > 1 else 1
+
+    is_scale = p.ndim == 1 or (p.ndim == 2 and (dim0 == 1 or dim1 == 1))
+    if is_scale:
+        return False, False, True
+
+    B = dim0 > dim1
+    A = dim0 < dim1
+    return B, A, False
+
+def apply_stiefel_update(
+    self,
+    p: Tensor,
+    group: Dict[str, Any],
+    update: Tensor,
+    lr: float | Tensor,
+    wd: float | None = None,
+    random_int_tensor: Tensor | None = None,
+    is_B: bool | None = None,
+    is_A: bool | None = None,
+    is_scale: bool | None = False,
+) -> None:
+    from ..util.param_update import _copy_stochastic_core_, copy_stochastic_
+    wd = group["weight_decay"] if wd is None else wd
+    cautious = group.get('cautious_wd', False)
+
+    if is_B or p.ndim == 1 or is_scale:
+        # Disable weight decay for the ortho matrix B or DoRA norm
+        wd = 0
+
+    if is_A:
+        # For matrix A, normalize weight decay by rank to make it invariant
+        wd = wd / p.shape[0]
+
+    scaled_wd = wd * (lr / self._init_lr)
+
+    # Compute full update in float32 if using bfloat16 with stochastic rounding
+    if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+        p_fp32 = p.float()
+        update_fp32 = update.float()
+
+        if is_B:
+            update_fp32 = trangent_proj(p_fp32, update_fp32, lr)
+            p_fp32.add_(-update_fp32)
+            p_fp32 = qr_retraction(p_fp32)
+            if random_int_tensor is not None:
+                _copy_stochastic_core_(p, p_fp32, random_int_tensor)
+                del random_int_tensor
+            else:
+                copy_stochastic_(p, p_fp32)
+            del p_fp32, update_fp32
+            return
+
+        # Apply weight decay if needed
+        if wd != 0:
+            if cautious:
+                # Cautious Weight Decay
+                mask = (update_fp32 * p_fp32 >= 0).float()
+                p_fp32.addcmul_(p_fp32, mask, value=-scaled_wd)
+                del mask
+            else:
+                # Standard decoupled weight decay
+                p_fp32.add_(p_fp32, alpha=-scaled_wd)
+
+        # Apply main update
+        p_fp32.add_(-update_fp32)
+
+        # Single stochastic rounding at the end
+        if random_int_tensor is not None:
+            # Compiled path: use the pre-computed random tensor
+            _copy_stochastic_core_(p, p_fp32, random_int_tensor)
+            del random_int_tensor
+        else:
+            # Uncompiled path: generate randoms inside
+            copy_stochastic_(p, p_fp32)
+        del p_fp32, update_fp32
+
+    else:
+        # Standard path for non-bfloat16 or without stochastic rounding
+
+        if is_B:
+            update = trangent_proj(p, update, lr)
+            p.add_(-update)
+            p = qr_retraction(p)
+            return
+
+        if wd != 0:
+            if cautious:
+                # Cautious Weight Decay
+                mask = (update * p >= 0).to(p.dtype)
+                p.addcmul_(p, mask, value=-scaled_wd)
+                del mask
+            else:
+                # Standard decoupled weight decay
+                p.add_(p, alpha=-scaled_wd)
+
+        # Apply main update
+        p.add_(-update)
+
+    del update

{adv_optm-2.3.dev1 → adv_optm-2.3.dev3}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-2.3.dev1 → adv_optm-2.3.dev3}/adv_optm.egg-info/SOURCES.txt

@@ -17,6 +17,9 @@ adv_optm/optim/Prodigy_adv.py
 adv_optm/optim/SignSGD_adv.py
 adv_optm/optim/Simplified_AdEMAMix.py
 adv_optm/optim/__init__.py
+adv_optm/stiefel_optm/Stiefel_LoRA.py
+adv_optm/stiefel_optm/__init__.py
+adv_optm/stiefel_optm/stiefel_util.py
 adv_optm/util/Kourkoutas.py
 adv_optm/util/Muon_AuxAdam.py
 adv_optm/util/Muon_util.py

{adv_optm-2.3.dev1 → adv_optm-2.3.dev3}/setup.py

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