adv-optm 1.0.6__py3-none-any.whl → 1.1.0.dev2__py3-none-any.whl

adv_optm/__init__.py

@@ -16,4 +16,4 @@ __all__ = [
     "Lion_Prodigy_adv",
 ]
 
-__version__ = "1.0.6"
+__version__ = "1.1.0.dev2"

adv_optm/optim/AdamW_adv.py

@@ -1,11 +1,12 @@
 import torch
-from typing import Optional
+from typing import Optional, Callable
 
 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
+from ..util.Kourkoutas import KourkoutasHelper
 
 class AdamW_adv(torch.optim.Optimizer):
     """
@@ -54,6 +55,28 @@ class AdamW_adv(torch.optim.Optimizer):
             as it gradually introduces the stabilizing slow momentum term. During
             the warmup, `alpha` ramps from 0 to its target value. If `None`,
             the scheduler is disabled. (default: None)
+        kourkoutas_beta (bool): whether to enable the layer-wise dynamic β₂ logic.
+            If `False`, the optimizer behaves as standard AdamW. (default: False)
+        beta2_min (float): The minimum value for dynamic β₂, used during periods of
+            high gradient variance ("sunspikes"). Must be less than `betas[1]`.
+            (default: 0.88)
+        ema_alpha (float): The decay rate for the Exponential Moving Average (EMA) of
+            the pooled gradient norms. Corresponds to `α` in the paper.
+            (default: 0.93)
+        tiny_spike (float): A small constant added to the denominator of the
+            "sunspike" ratio calculation to prevent division by zero. Corresponds
+            to `ε_spike` in the paper. (default: 1e-9)
+        k_warmup_steps (int): The number of initial steps during which β₂ is held
+            at a fixed average value (`(beta2_min + beta2_max) / 2`) before the
+            dynamic logic activates. (default: 0)
+        k_logging (int): if > 0 and kourkoutas_beta=True, enables periodic console
+            logging of Kourkoutas-β statistics (min, max, mean of `β₂` across layers)
+            every logging steps. Useful for debugging and tuning. Set to 0 to disable
+            logging (default: 0). 
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
+            and returns a unique, hashable key representing its "layer" or "bucket".
+            If `None`, parameters are bucketed by their memory ID (tensor-wise).
+            (default: None)
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
     """
@@ -76,6 +99,13 @@ class AdamW_adv(torch.optim.Optimizer):
         beta3_ema: float = 0.9999,
         alpha: float = 5.0,
         t_alpha: int | None = None,
+        kourkoutas_beta: bool = False,
+        beta2_min: float = 0.88,
+        ema_alpha: float = 0.93,
+        tiny_spike: float = 1e-9,
+        k_warmup_steps: int = 0,
+        k_logging: int = 0,
+        layer_key_fn: Optional[Callable] = None,
         nnmf_factor: bool = False,
     ):
         if not (lr >= 0.0):
@@ -86,6 +116,8 @@ class AdamW_adv(torch.optim.Optimizer):
             raise ValueError(f"Epsilon should be >= 0.0. Got {eps}")
         if not (weight_decay >= 0.0):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
+        if kourkoutas_beta and not (betas[1] > beta2_min): raise ValueError(f"For Kourkoutas-β, betas[1] (as beta2_max) must be > beta2_min. Got {betas[1]} and {beta2_min}")
+
         if cautious_mask and grams_moment:
             print("Warning: cautious is incompatible with grams, Disabling cautious.")
             cautious_mask = False
@@ -95,14 +127,21 @@ class AdamW_adv(torch.optim.Optimizer):
             "vector_reshape": vector_reshape, "use_atan2": use_atan2,
             "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction,
             "beta3_ema": beta3_ema, "alpha": alpha, "t_alpha": t_alpha,
+            "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
+            "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
         }
         self.stochastic_rounding = stochastic_rounding
         self.cautious_mask = cautious_mask
         self.grams_moment = grams_moment
         self.use_AdEMAMix = use_AdEMAMix
         self.factored = nnmf_factor
+        self.kourkoutas_beta = kourkoutas_beta
+        self.layer_key_fn = layer_key_fn
         super().__init__(params, defaults)
 
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
+
     @property
     def supports_fused_back_pass(self):
         return True
@@ -127,8 +166,6 @@ 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
@@ -148,7 +185,7 @@ class AdamW_adv(torch.optim.Optimizer):
                 d1, d2 = state['effective_shape']
 
                 # First moment (m)
-                if beta1 > 0:
+                if group['betas'][0] > 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.grams_moment:
@@ -163,16 +200,31 @@ 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
-                if beta1 > 0:
+                if group['betas'][0] > 0:
                     state['exp_avg'] = torch.zeros_like(p, device=device, dtype=dtype)
                 if self.use_AdEMAMix:
                     state['exp_avg_slow'] = torch.zeros_like(p, device=device, dtype=dtype)
                 state['exp_avg_sq'] = torch.zeros_like(p, device=device, dtype=dtype)
 
+        beta1, beta2 = group['betas']
+
+        current_step = state['step']
+        if group['kourkoutas_beta']:
+            # Call prepare_step() once at the beginning of the step for all params
+            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            # Accumulate current grad's norm for the *next* step
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+            # Get the dynamic beta2 calculated in prepare_step()
+            beta2 = self.kourkoutas_helper.get_beta2(p, group, current_step)
+
         step = state['step'] + 1
         if group['use_bias_correction']:
             bias_correction1 = 1.0 - beta1 ** step
-            bias_correction2 = 1.0 - beta2 ** step
+            if group['kourkoutas_beta']:
+                bias_correction2 = 1.0 - group['betas'][1] ** step
+                # Use beta2_max for bias correction
+            else:
+                bias_correction2 = 1.0 - beta2 ** step
         else:
             bias_correction1 = 1
             bias_correction2 = 1
@@ -315,4 +367,4 @@ class AdamW_adv(torch.optim.Optimizer):
             for i, p in enumerate(group['params']):
                 self.step_parameter(p, group, i)
 
-        return loss
+        return loss

adv_optm/optim/Adopt_adv.py

@@ -6,6 +6,7 @@ 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
+from ..util.Kourkoutas import KourkoutasHelper
 
 class Adopt_adv(torch.optim.Optimizer):
     """
@@ -72,6 +73,28 @@ class Adopt_adv(torch.optim.Optimizer):
             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)
+        kourkoutas_beta (bool): whether to enable the layer-wise dynamic β₂ logic.
+            If `False`, the optimizer behaves as standard Adopt. (default: False)
+        beta2_min (float): The minimum value for dynamic β₂, used during periods of
+            high gradient variance ("sunspikes"). Must be less than `betas[1]`.
+            (default: 0.88)
+        ema_alpha (float): The decay rate for the Exponential Moving Average (EMA) of
+            the pooled gradient norms. Corresponds to `α` in the paper.
+            (default: 0.93)
+        tiny_spike (float): A small constant added to the denominator of the
+            "sunspike" ratio calculation to prevent division by zero. Corresponds
+            to `ε_spike` in the paper. (default: 1e-9)
+        k_warmup_steps (int): The number of initial steps during which β₂ is held
+            at a fixed average value (`(beta2_min + beta2_max) / 2`) before the
+            dynamic logic activates. (default: 0)
+        k_logging (int): if > 0 and kourkoutas_beta=True, enables periodic console
+            logging of Kourkoutas-β statistics (min, max, mean of `β₂` across layers)
+            every logging steps. Useful for debugging and tuning. Set to 0 to disable
+            logging (default: 0). 
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
+            and returns a unique, hashable key representing its "layer" or "bucket".
+            If `None`, parameters are bucketed by their memory ID (tensor-wise).
+            (default: None)
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
     """
@@ -96,6 +119,13 @@ class Adopt_adv(torch.optim.Optimizer):
         t_alpha: int | None = None,
         Simplified_AdEMAMix: bool = False,
         alpha_grad: float = 100.0,
+        kourkoutas_beta: bool = False,
+        beta2_min: float = 0.88,
+        ema_alpha: float = 0.93,
+        tiny_spike: float = 1e-9,
+        k_warmup_steps: int = 0,
+        k_logging: int = 0,
+        layer_key_fn: Optional[Callable] = None,
         nnmf_factor: bool = False,
     ):
         if not (lr >= 0.0):
@@ -111,6 +141,7 @@ class Adopt_adv(torch.optim.Optimizer):
             cautious_mask = False
         if betas[0] == 0.0 and Simplified_AdEMAMix:
             raise ValueError(f"Beta1 cannot be 0.0 when using Simplified_AdEMAMix. Got {betas[0]}")
+        if kourkoutas_beta and not (betas[1] > beta2_min): raise ValueError(f"For Kourkoutas-β, betas[1] (as beta2_max) must be > beta2_min. Got {betas[1]} and {beta2_min}")
         if use_AdEMAMix and Simplified_AdEMAMix:
             print("Warning: use_AdEMAMix is incompatible with Simplified_AdEMAMix, Disabling use_AdEMAMix.")
         if grams_moment and Simplified_AdEMAMix:
@@ -125,6 +156,8 @@ class Adopt_adv(torch.optim.Optimizer):
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
             "vector_reshape": vector_reshape, "beta3_ema": beta3_ema, "alpha": alpha,
             "t_alpha": t_alpha, "alpha_grad": alpha_grad, 
+            "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
+            "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
         }
         self.clip_lambda = clip_lambda
         self.stochastic_rounding = stochastic_rounding
@@ -135,8 +168,13 @@ class Adopt_adv(torch.optim.Optimizer):
         self.use_AdEMAMix = use_AdEMAMix and not Simplified_AdEMAMix
         self.Simplified_AdEMAMix = Simplified_AdEMAMix
         self.factored = nnmf_factor
+        self.kourkoutas_beta = kourkoutas_beta
+        self.layer_key_fn = layer_key_fn
         super().__init__(params, defaults)
 
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
+
     @property
     def supports_fused_back_pass(self): return True
     @property
@@ -156,8 +194,6 @@ class Adopt_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
@@ -176,7 +212,7 @@ class Adopt_adv(torch.optim.Optimizer):
                 d1, d2 = state['effective_shape']
 
                 # m_0 = 0
-                if beta1 > 0:
+                if group['betas'][0] > 0:
                     state['mu_m_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype) 
                     state['mv_m_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
                     if not self.grams_moment:
@@ -195,12 +231,23 @@ class Adopt_adv(torch.optim.Optimizer):
                 # Initialize v_0 using NMF
                 _nnmf(vt_init, out=(state['mu_v_nmf'], state['mv_v_nmf']))
             else: # Fallback for non-factored tensors
-                if beta1 > 0:
+                if group['betas'][0] > 0:
                     state['exp_avg'] = torch.zeros_like(p, dtype=dtype) # m_0
                 if self.use_AdEMAMix:
                     state['exp_avg_slow'] = torch.zeros_like(p, dtype=dtype)
                 state['exp_avg_sq'] = grad.square()   # v_0
 
+        beta1, beta2 = group['betas']
+
+        current_step = state['step']
+        if group['kourkoutas_beta']:
+            # Call prepare_step() once at the beginning of the step for all params
+            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            # Accumulate current grad's norm for the *next* step
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+            # Get the dynamic beta2 calculated in prepare_step()
+            beta2 = self.kourkoutas_helper.get_beta2(p, group, current_step)
+
         # The first step is for initialization only (skip when use_atan2 as it's scale invariant).
         if state['step'] == 0 and not self.use_atan2:
             state['step'] += 1
@@ -211,10 +258,10 @@ class Adopt_adv(torch.optim.Optimizer):
             alpha = group['alpha']
             t_alpha = group['t_alpha']
             # Use step+1 for 1-based step count in scheduler
-            current_step = state['step'] + 1
+            alpha_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 t_alpha is not None and t_alpha > 0 and alpha_step < t_alpha:
+                alpha_t = min(alpha_step * alpha / t_alpha, alpha)
         if self.Simplified_AdEMAMix:
             alpha_grad = group["alpha_grad"]
 
@@ -386,4 +433,8 @@ class Adopt_adv(torch.optim.Optimizer):
             for i, p in enumerate(group['params']):
                 self.step_parameter(p, group, i)
 
+        if self.kourkoutas_beta and self.k_logging > 0 and hasattr(self, '_beta2_log'):
+            first_param_state = self.state[self.param_groups[0]['params'][0]]
+            step_num = first_param_state['step']
+
         return loss

adv_optm/optim/Prodigy_adv.py

@@ -3,11 +3,14 @@ import torch.distributed as dist
 
 import math
 
+from typing import Optional, Callable
+
 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
+from ..util.Kourkoutas import KourkoutasHelper
 
 class Prodigy_adv(torch.optim.Optimizer):
     """
@@ -85,6 +88,28 @@ class Prodigy_adv(torch.optim.Optimizer):
         prodigy_steps (int): If greater than zero, disable Prodigy's stepsize adjustments
             after the specified optimiser step and release all state memory required by Prodigy
             (default: 0).
+        kourkoutas_beta (bool): whether to enable the layer-wise dynamic β₂ logic.
+            If `False`, the optimizer behaves as standard AdamW/Prodigy. (default: False)
+        beta2_min (float): The minimum value for dynamic β₂, used during periods of
+            high gradient variance ("sunspikes"). Must be less than `betas[1]`.
+            (default: 0.88)
+        ema_alpha (float): The decay rate for the Exponential Moving Average (EMA) of
+            the pooled gradient norms. Corresponds to `α` in the paper.
+            (default: 0.93)
+        tiny_spike (float): A small constant added to the denominator of the
+            "sunspike" ratio calculation to prevent division by zero. Corresponds
+            to `ε_spike` in the paper. (default: 1e-9)
+        k_warmup_steps (int): The number of initial steps during which β₂ is held
+            at a fixed average value (`(beta2_min + beta2_max) / 2`) before the
+            dynamic logic activates. (default: 0)
+        k_logging (int): if > 0 and kourkoutas_beta=True, enables periodic console
+            logging of Kourkoutas-β statistics (min, max, mean of `β₂` across layers)
+            every logging steps. Useful for debugging and tuning. Set to 0 to disable
+            logging (default: 0). 
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
+            and returns a unique, hashable key representing its "layer" or "bucket".
+            If `None`, parameters are bucketed by their memory ID (tensor-wise).
+            (default: None)
     """
 
     def __init__(
@@ -116,6 +141,13 @@ class Prodigy_adv(torch.optim.Optimizer):
         fsdp_in_use: bool = False,
         slice_p: int = 11,
         prodigy_steps: int = 0,
+        kourkoutas_beta: bool = False,
+        beta2_min: float = 0.88,
+        ema_alpha: float = 0.93,
+        tiny_spike: float = 1e-9,
+        k_warmup_steps: int = 0,
+        k_logging: int = 0,
+        layer_key_fn: Optional[Callable] = None,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -141,6 +173,8 @@ class Prodigy_adv(torch.optim.Optimizer):
         if use_atan2 and Simplified_AdEMAMix:
             print("Warning: use_atan2 is incompatible with Simplified_AdEMAMix. Disabling use_atan2.")
             use_atan2 = False
+        if kourkoutas_beta and not (betas[1] > beta2_min):
+            raise ValueError(f"For Kourkoutas-β, betas[1] (as beta2_max) must be > beta2_min. Got {betas[1]} and {beta2_min}")
         if Simplified_AdEMAMix and alpha_grad > 0:
             # scales d_coef by alpha_grad, this force prodigy to behave well with Simplified_AdEMAMix
             d_coef = d_coef/alpha_grad
@@ -153,7 +187,9 @@ class Prodigy_adv(torch.optim.Optimizer):
             "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, "prodigy_steps": prodigy_steps,
-            "alpha_grad": alpha_grad, 
+            "alpha_grad": alpha_grad,
+            "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
+            "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
         }
         self.stochastic_rounding = stochastic_rounding
         self.cautious_mask = cautious_mask and not Simplified_AdEMAMix
@@ -162,7 +198,13 @@ class Prodigy_adv(torch.optim.Optimizer):
         self.Simplified_AdEMAMix = Simplified_AdEMAMix
         self.factored = nnmf_factor
         self.fsdp_in_use = fsdp_in_use
+        
+        self.kourkoutas_beta = kourkoutas_beta
+        self.layer_key_fn = layer_key_fn
+
         super().__init__(params, defaults)
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
         self.init_step()
 
     @property
@@ -180,19 +222,17 @@ class Prodigy_adv(torch.optim.Optimizer):
     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.beta1, self.beta2_default = g_group['betas']
         self.beta3 = g_group['beta3']
         if self.beta3 is None:
-            self.beta3 = math.sqrt(self.beta2)
+            self.beta3 = math.sqrt(self.beta2_default)
         
-        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()
@@ -258,14 +298,25 @@ class Prodigy_adv(torch.optim.Optimizer):
             else:
                 state['p0'] = torch.tensor(0, device=device, dtype=p.dtype)
 
+        current_step = state['step']
+        if group['kourkoutas_beta']:
+            # Call prepare_step() once at the beginning of the step for all params
+            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            # Accumulate current grad's norm for the *next* step
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+            # Get the dynamic beta2 calculated in prepare_step()
+            beta2 = self.kourkoutas_helper.get_beta2(p, group, current_step)
+        else:
+            beta2 = self.beta2_default
+
         if self.use_AdEMAMix:
             beta3_ema = group['beta3_ema']
             alpha = group['alpha']
             t_alpha = group['t_alpha']
-            current_step = state['step'] + 1
+            alpha_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 t_alpha is not None and t_alpha > 0 and alpha_step < t_alpha:
+                alpha_t = min(alpha_step * alpha / t_alpha, alpha)
         if self.Simplified_AdEMAMix:
             alpha_grad = group["alpha_grad"]
 
@@ -295,7 +346,7 @@ class Prodigy_adv(torch.optim.Optimizer):
                     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))
+            vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=self.d * self.d * (1.0 - beta2))
 
             if self.use_AdEMAMix:
                 mt_slow = _unnmf((state['mu_m_slow_nmf'], state['mv_m_slow_nmf']))
@@ -368,7 +419,7 @@ class Prodigy_adv(torch.optim.Optimizer):
             else:
                 update = exp_avg.clone() if self.beta1 > 0 else grad.mul(self.d)
 
-            exp_avg_sq.mul_(self.beta2).addcmul_(grad, grad.conj(), value=self.d * self.d * (1.0 - self.beta2))
+            exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=self.d * self.d * (1.0 - beta2))
 
             if group['use_atan2']:
                 a = 1.2732395
@@ -431,7 +482,6 @@ class Prodigy_adv(torch.optim.Optimizer):
             for i, p in enumerate(group['params']):
                 self.step_parameter(p, group, i)
 
-
         self.calculate_d()
         self.init_step()
         return loss

adv_optm/optim/Simplified_AdEMAMix.py

@@ -1,4 +1,5 @@
 import torch
+from typing import Optional, Callable
 
 import math
 
@@ -7,6 +8,7 @@ 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
+from ..util.Kourkoutas import KourkoutasHelper
 
 # A little helper from the original simplified_AdEMAMix
 def linear_hl_warmup_scheduler(step, beta_end, beta_start=0, warmup=1):
@@ -47,6 +49,28 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
         stochastic_rounding (bool): whether to use stochastic
             rounding for BF16 parameter updates (default: True).
         orthogonal_gradient (bool): whether to use OrthoGrad. (default: False)
+        kourkoutas_beta (bool): whether to enable the layer-wise dynamic β₂ logic.
+            If `False`, the optimizer behaves as standard Simplified_AdEMAMix. (default: False)
+        beta2_min (float): The minimum value for dynamic β₂, used during periods of
+            high gradient variance ("sunspikes"). Must be less than `betas[1]`.
+            (default: 0.88)
+        ema_alpha (float): The decay rate for the Exponential Moving Average (EMA) of
+            the pooled gradient norms. Corresponds to `α` in the paper.
+            (default: 0.93)
+        tiny_spike (float): A small constant added to the denominator of the
+            "sunspike" ratio calculation to prevent division by zero. Corresponds
+            to `ε_spike` in the paper. (default: 1e-9)
+        k_warmup_steps (int): The number of initial steps during which β₂ is held
+            at a fixed average value (`(beta2_min + beta2_max) / 2`) before the
+            dynamic logic activates. (default: 0)
+        k_logging (int): if > 0 and kourkoutas_beta=True, enables periodic console
+            logging of Kourkoutas-β statistics (min, max, mean of `β₂` across layers)
+            every logging steps. Useful for debugging and tuning. Set to 0 to disable
+            logging (default: 0). 
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
+            and returns a unique, hashable key representing its "layer" or "bucket".
+            If `None`, parameters are bucketed by their memory ID (tensor-wise).
+            (default: None)
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
     """
@@ -65,6 +89,13 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
         vector_reshape: bool = True,
         stochastic_rounding: bool = True,
         orthogonal_gradient: bool = False,
+        kourkoutas_beta: bool = False,
+        beta2_min: float = 0.88,
+        ema_alpha: float = 0.93,
+        tiny_spike: float = 1e-9,
+        k_warmup_steps: int = 0,
+        k_logging: int = 0,
+        layer_key_fn: Optional[Callable] = None,
         nnmf_factor: bool = False,
     ):
         if not (lr >= 0.0):
@@ -77,17 +108,25 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
         if not 0.0 <= alpha_grad:
             raise ValueError("Invalid alpha value: {}".format(alpha_grad))
+        if kourkoutas_beta and not (betas[1] > beta2_min): raise ValueError(f"For Kourkoutas-β, betas[1] (as beta2_max) must be > beta2_min. Got {betas[1]} and {beta2_min}")
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
             "alpha_grad": alpha_grad, "beta1_warmup": beta1_warmup, "min_beta1": min_beta1,
             "vector_reshape": vector_reshape,
             "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction,
+            "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
+            "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
         }
         self.stochastic_rounding = stochastic_rounding
         self.factored = nnmf_factor
+        self.kourkoutas_beta = kourkoutas_beta
+        self.layer_key_fn = layer_key_fn
         super().__init__(params, defaults)
 
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
+
     @property
     def supports_fused_back_pass(self):
         return True
@@ -150,6 +189,16 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
                 state['den_sum'] = 1.0
 
         beta1_final, beta2 = group["betas"]
+
+        current_step = state['step']
+        if group['kourkoutas_beta']:
+            # Call prepare_step() once at the beginning of the step for all params
+            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            # Accumulate current grad's norm for the *next* step
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+            # Get the dynamic beta2 calculated in prepare_step()
+            beta2 = self.kourkoutas_helper.get_beta2(p, group, current_step)
+
         beta1_warmup = group["beta1_warmup"]
         alpha_grad = group["alpha_grad"]
 
@@ -161,7 +210,10 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
 
         if group['use_bias_correction']:
             state['num_sum'] = beta1 * state['num_sum'] + 1.0
-            state['den_sum'] = beta2 * state['den_sum'] + (1.0 - beta2)
+            if group['kourkoutas_beta']:
+                state['den_sum'] = group['betas'][1] * state['den_sum'] + (1.0 - group['betas'][1])
+            else:
+                state['den_sum'] = beta2 * state['den_sum'] + (1.0 - beta2)
 
         if state['factored']:
             d1, d2 = state['effective_shape']

adv_optm/util/Kourkoutas.py

@@ -0,0 +1,134 @@
+import torch
+from torch.optim import Optimizer
+from typing import Callable
+
+class KourkoutasHelper:
+    """
+    A helper class to add layer-wise Kourkoutas-β functionality to a PyTorch optimizer.
+    """
+    def __init__(self, optimizer: Optimizer):
+        # We need a reference to the optimizer to access its param_groups and state
+        if not hasattr(optimizer, 'param_groups'):
+            raise TypeError("optimizer must be a valid torch.optim.Optimizer instance.")
+        self.optimizer = optimizer
+        
+        # State managed by the helper
+        self.layer_state = {}
+        self.layer_info = {}
+        self._layer_info_built = False
+        self._current_step_prepared = -1
+
+    def _build_layer_info_if_needed(self):
+        """Builds a map of layers and the parameters they contain."""
+        if self._layer_info_built:
+            return
+            
+        if not hasattr(self.optimizer, 'layer_key_fn') or self.optimizer.layer_key_fn is None:
+            print("Warning: KourkoutasHelper requires 'layer_key_fn' on the optimizer. Defaulting to tensor-wise (id).")
+            self.optimizer.layer_key_fn = lambda p: id(p)
+
+        for group in self.optimizer.param_groups:
+            for p in group['params']:
+                if p.grad is None: continue
+                layer_key = self.optimizer.layer_key_fn(p)
+                if layer_key not in self.layer_info:
+                    self.layer_info[layer_key] = {'params': [], 'group_ref': group}
+                self.layer_info[layer_key]['params'].append(p)
+        
+        k_logging_interval = self.optimizer.param_groups[0].get('k_logging', 0)
+        if k_logging_interval > 0:
+            print(f"[Kourkoutas-β Debug] Layer info built. Found {len(self.layer_info)} unique layers/buckets.")
+
+        self._layer_info_built = True
+
+    def prepare_step(self, current_step: int):
+        """
+        Calculates dynamic beta2 for all layers using the completed scalar accumulators
+        from the PREVIOUS step. Should be called once at the start of an optimizer step.
+        """
+        self._build_layer_info_if_needed()
+
+        # Check if logging is enabled for this step based on the interval
+        k_logging_interval = self.optimizer.param_groups[0].get('k_logging', 0)
+        is_logging_step = k_logging_interval > 0 and (current_step + 1) % k_logging_interval == 0
+
+        beta2_log = [] if is_logging_step else None
+        first_layer_key = next(iter(self.layer_info), None)
+
+        for layer_key, info in self.layer_info.items():
+            params, group = info['params'], info['group_ref']
+            
+            if layer_key not in self.layer_state:
+                self.layer_state[layer_key] = {
+                    'r_ema_grad_norm': torch.tensor(0.0, device=params[0].device, dtype=torch.float32),
+                    'sum_sq_accumulator': torch.tensor(0.0, device=params[0].device, dtype=torch.float32)
+                }
+            
+            layer_state = self.layer_state[layer_key]
+            
+            # Use the completed accumulator from the previous step
+            pooled_grad_norm = torch.sqrt(layer_state['sum_sq_accumulator'])
+            
+            r_ema = layer_state['r_ema_grad_norm']
+            prev_r_ema_val = r_ema.item() # for logging
+            
+            # EMA is always updated, even during warmup
+            r_ema.mul_(group['ema_alpha']).add_(pooled_grad_norm, alpha=1.0 - group['ema_alpha'])
+            
+            sun = torch.tensor(0.0, device=r_ema.device) # Default sun to 0 for warmup
+            beta2_max = group['betas'][1]
+
+            # --- CONSOLIDATED WARMUP LOGIC ---
+            if current_step < group['k_warmup_steps']:
+                beta2 = beta2_max
+            else:
+                raw = pooled_grad_norm / (r_ema + group['tiny_spike'])
+                sun = raw / (1.0 + raw)
+                beta2 = beta2_max - (beta2_max - group['beta2_min']) * sun
+
+            layer_state['dynamic_beta2'] = beta2.item() if isinstance(beta2, torch.Tensor) else beta2
+            layer_state['sum_sq_accumulator'].zero_()
+
+            if is_logging_step:
+                beta2_log.append(layer_state['dynamic_beta2'])
+                if layer_key == first_layer_key:
+                    print(f"\n[Kourkoutas-β Debug] Step {current_step + 1} - Sample Layer '{layer_key}':")
+                    print(f"  - Grad Norm: {pooled_grad_norm.item():.4e}, Prev EMA: {prev_r_ema_val:.4e}, New EMA: {r_ema.item():.4e}")
+                    print(f"  - Sunspike: {sun.item():.4f}, Dynamic Beta2: {layer_state['dynamic_beta2']:.4f}")
+        
+        if is_logging_step and beta2_log:
+            beta2_tensor = torch.tensor(beta2_log, device='cpu')
+            print(f"[Kourkoutas-β Debug] Step {current_step + 1} Overall Beta2 Stats: Min={beta2_tensor.min():.4f}, Max={beta2_tensor.max():.4f}, Mean={beta2_tensor.mean():.4f}")
+
+
+    def maybe_prepare_step(self, current_step: int):
+        """
+        A universal guard that calls prepare_step() exactly once per training step.
+        """
+        if self._current_step_prepared < current_step:
+            self.prepare_step(current_step)
+            self._current_step_prepared = current_step
+
+    def accumulate_gradient_sq_norm(self, p: torch.Tensor, grad: torch.Tensor):
+        """
+        Accumulates the squared L2 norm of a single gradient for the next step's calculation.
+        """
+        self._build_layer_info_if_needed()
+        layer_key = self.optimizer.layer_key_fn(p)
+
+        if layer_key in self.layer_info:
+            if layer_key not in self.layer_state:
+                    self.layer_state[layer_key] = {
+                    'r_ema_grad_norm': torch.tensor(0.0, device=p.device, dtype=torch.float32),
+                    'sum_sq_accumulator': torch.tensor(0.0, device=p.device, dtype=torch.float32)
+                }
+            # Accumulate for the *next* step's prepare_step call
+            self.layer_state[layer_key]['sum_sq_accumulator'] += torch.sum(grad.detach().pow(2)).float()
+
+    def get_beta2(self, p: torch.Tensor, group: dict, current_step: int) -> float:
+        """
+        Gets the appropriate beta2 for the current parameter, handling warmup and dynamic value fetching.
+        """
+        layer_key = self.optimizer.layer_key_fn(p)
+        # The default is the max value, which is correct for unmapped params or edge cases
+        return self.layer_state.get(layer_key, {}).get('dynamic_beta2', group['betas'][1])

{adv_optm-1.0.6.dist-info → adv_optm-1.1.0.dev2.dist-info}/METADATA

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

adv_optm-1.1.0.dev2.dist-info/RECORD

@@ -0,0 +1,20 @@
+adv_optm/__init__.py,sha256=hkmbLr1AVDoC6VbnyTkNy-G4g5bmcLFH2Kv4dYWp9uY,311
+adv_optm/optim/AdamW_adv.py,sha256=H4XlYZELwiFvXt0A9wMlRNiw9c8rmPMspHDCvR_SZIQ,17487
+adv_optm/optim/Adopt_adv.py,sha256=PJ3ZaLgzYbvxXDS56FGjzMrVMyHDXSWdUPHnX5NpNAA,21241
+adv_optm/optim/Lion_Prodigy_adv.py,sha256=sGzhts9a6gHfCkuHTB5L9IrClo4c6UThzYYErBwqOaA,12844
+adv_optm/optim/Lion_adv.py,sha256=6G1CukJB_pC7l9HwFEuY1ydsNHZFabVmOvcHDsHHVuQ,8295
+adv_optm/optim/Prodigy_adv.py,sha256=-eMTutexbGrUQtSehKaOo6BO_p3QySpSIMgJKWvbxog,25517
+adv_optm/optim/Simplified_AdEMAMix.py,sha256=b4GaSI-TX6wFBqGxZeoJPbf2nVRCEtB3WVb1olDgY14,12980
+adv_optm/optim/__init__.py,sha256=pcP865H2j1tut2VfTUhzQh7V8TF_tzPjqFnjMfFed2k,382
+adv_optm/util/BF16_Stochastic_Rounding.py,sha256=Q5H0BcogmE4atP65dLoI21HKSf50lRdsBDfeF6v9Tbg,1548
+adv_optm/util/Effective_Shape.py,sha256=TBvIk1V8IuTbbBsxuekJA4e_v8JlR5Nujtut8RTWAm4,318
+adv_optm/util/Kourkoutas.py,sha256=6OzK96KJ7Dd9Py8hiGWszF9C_n4uVoDjFCA_EYbhL4c,6600
+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-1.1.0.dev2.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-1.1.0.dev2.dist-info/METADATA,sha256=Y2F2wkpPmdbRtHft1KdCm1D6feTmiP5kFJ6iYpSLwCo,8427
+adv_optm-1.1.0.dev2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+adv_optm-1.1.0.dev2.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-1.1.0.dev2.dist-info/RECORD,,

adv_optm-1.0.6.dist-info/RECORD

@@ -1,19 +0,0 @@
-adv_optm/__init__.py,sha256=dAbueuVEIGoYrYXx8UE4ATfFBH5wEKrpkXGPTjFH0r0,306
-adv_optm/optim/AdamW_adv.py,sha256=aTuYcJgd_EcZOrs6TDgBrBKw3wtU5LPzE5WvTBDDeEo,14317
-adv_optm/optim/Adopt_adv.py,sha256=FTpDDSlYruZDt1VVLgEI_bADiO8f26j-utQs7Gn2fFA,18108
-adv_optm/optim/Lion_Prodigy_adv.py,sha256=sGzhts9a6gHfCkuHTB5L9IrClo4c6UThzYYErBwqOaA,12844
-adv_optm/optim/Lion_adv.py,sha256=6G1CukJB_pC7l9HwFEuY1ydsNHZFabVmOvcHDsHHVuQ,8295
-adv_optm/optim/Prodigy_adv.py,sha256=G8xXLO9YBeLb9574uS0HpdY9w3ojblaV-PJFghUnToQ,22493
-adv_optm/optim/Simplified_AdEMAMix.py,sha256=tb3d6Cw_nGwcTzYUhDnKqyP7GzjD1hn8k4WqGG5lhmw,9813
-adv_optm/optim/__init__.py,sha256=pcP865H2j1tut2VfTUhzQh7V8TF_tzPjqFnjMfFed2k,382
-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-1.0.6.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-1.0.6.dist-info/METADATA,sha256=3PslWXH0ysoiXU83vN3F9kWRw48fwUM4H1z1tMyEGvI,8422
-adv_optm-1.0.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-1.0.6.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-1.0.6.dist-info/RECORD,,