adv-optm 1.0.6__tar.gz → 1.1.0.dev1__tar.gz

{adv_optm-1.0.6 → adv_optm-1.1.0.dev1}/PKG-INFO

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

{adv_optm-1.0.6 → adv_optm-1.1.0.dev1}/adv_optm/__init__.py

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

{adv_optm-1.0.6 → adv_optm-1.1.0.dev1}/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,6 +127,8 @@ 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,
         }
         self.stochastic_rounding = stochastic_rounding
         self.cautious_mask = cautious_mask
@@ -103,6 +137,12 @@ class AdamW_adv(torch.optim.Optimizer):
         self.factored = nnmf_factor
         super().__init__(params, defaults)
 
+        self.kourkoutas_beta = kourkoutas_beta
+        self.k_logging= k_logging and kourkoutas_beta
+        self.layer_key_fn = layer_key_fn and kourkoutas_beta
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
+
     @property
     def supports_fused_back_pass(self):
         return True
@@ -127,8 +167,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 +186,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 +201,29 @@ 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)
 
+        current_step = state['step']
+        if group['kourkoutas_beta']:
+            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+        
+        beta1, beta2 = group['betas']
+        if group['kourkoutas_beta']:
+            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 +366,14 @@ class AdamW_adv(torch.optim.Optimizer):
             for i, p in enumerate(group['params']):
                 self.step_parameter(p, group, i)
 
-        return loss
+        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']
+
+            if step_num > 0 and step_num % self.k_logging == 0:
+                if self._beta2_log:
+                    beta2_tensor = torch.tensor(self._beta2_log, device='cpu')
+                    print(f"Step {step_num}: Kourkoutas beta2 stats: Min={beta2_tensor.min():.4f}, Max={beta2_tensor.max():.4f}, Mean={beta2_tensor.mean():.4f}")
+                delattr(self, '_beta2_log')
+
+        return loss

{adv_optm-1.0.6 → adv_optm-1.1.0.dev1}/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,
         }
         self.clip_lambda = clip_lambda
         self.stochastic_rounding = stochastic_rounding
@@ -137,6 +170,12 @@ class Adopt_adv(torch.optim.Optimizer):
         self.factored = nnmf_factor
         super().__init__(params, defaults)
 
+        self.kourkoutas_beta = kourkoutas_beta
+        self.k_logging= k_logging and kourkoutas_beta
+        self.layer_key_fn = layer_key_fn and kourkoutas_beta
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
+
     @property
     def supports_fused_back_pass(self): return True
     @property
@@ -156,8 +195,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 +213,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 +232,21 @@ 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
 
+        current_step = state['step']
+        if group['kourkoutas_beta']:
+            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+
+        beta1, beta2 = group['betas']
+        if group['kourkoutas_beta']:
+            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
@@ -386,4 +432,14 @@ 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']
+
+            if step_num > 0 and step_num % self.k_logging == 0:
+                if self._beta2_log:
+                    beta2_tensor = torch.tensor(self._beta2_log, device='cpu')
+                    print(f"Step {step_num}: Kourkoutas beta2 stats: Min={beta2_tensor.min():.4f}, Max={beta2_tensor.max():.4f}, Mean={beta2_tensor.mean():.4f}")
+                delattr(self, '_beta2_log')
+
         return loss

{adv_optm-1.0.6 → adv_optm-1.1.0.dev1}/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,
         }
         self.stochastic_rounding = stochastic_rounding
         self.cautious_mask = cautious_mask and not Simplified_AdEMAMix
@@ -163,6 +199,13 @@ class Prodigy_adv(torch.optim.Optimizer):
         self.factored = nnmf_factor
         self.fsdp_in_use = fsdp_in_use
         super().__init__(params, defaults)
+
+        self.kourkoutas_beta = kourkoutas_beta
+        self.k_logging= k_logging and kourkoutas_beta
+        self.layer_key_fn = layer_key_fn and kourkoutas_beta
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
+
         self.init_step()
 
     @property
@@ -180,19 +223,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,6 +299,15 @@ 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']:
+            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+
+        beta2 = self.beta2_default
+        if group['kourkoutas_beta']:
+            beta2 = self.kourkoutas_helper.get_beta2(p, group, current_step)
+
         if self.use_AdEMAMix:
             beta3_ema = group['beta3_ema']
             alpha = group['alpha']
@@ -295,7 +345,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 +418,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,6 +481,15 @@ class Prodigy_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']
+
+            if step_num > 0 and step_num % self.k_logging == 0:
+                if self._beta2_log:
+                    beta2_tensor = torch.tensor(self._beta2_log, device='cpu')
+                    print(f"Step {step_num}: Kourkoutas beta2 stats: Min={beta2_tensor.min():.4f}, Max={beta2_tensor.max():.4f}, Mean={beta2_tensor.mean():.4f}")
+                delattr(self, '_beta2_log')
 
         self.calculate_d()
         self.init_step()

{adv_optm-1.0.6 → adv_optm-1.1.0.dev1}/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,26 @@ 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,
         }
         self.stochastic_rounding = stochastic_rounding
         self.factored = nnmf_factor
         super().__init__(params, defaults)
 
+        self.kourkoutas_beta = kourkoutas_beta
+        self.k_logging= k_logging and kourkoutas_beta
+        self.layer_key_fn = layer_key_fn and kourkoutas_beta
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
+
     @property
     def supports_fused_back_pass(self):
         return True
@@ -149,9 +189,17 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
                 state['num_sum'] = 1.0
                 state['den_sum'] = 1.0
 
+        current_step = state['step']
+        if group['kourkoutas_beta']:
+            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+
         beta1_final, beta2 = group["betas"]
         beta1_warmup = group["beta1_warmup"]
         alpha_grad = group["alpha_grad"]
+        
+        if group['kourkoutas_beta']:
+            beta2 = self.kourkoutas_helper.get_beta2(p, group, current_step)
 
         if beta1_warmup is not None:
             step = state['step'] + 1
@@ -161,7 +209,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']
@@ -243,4 +294,14 @@ class Simplified_AdEMAMix(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']
+
+            if step_num > 0 and step_num % self.k_logging == 0:
+                if self._beta2_log:
+                    beta2_tensor = torch.tensor(self._beta2_log, device='cpu')
+                    print(f"Step {step_num}: Kourkoutas beta2 stats: Min={beta2_tensor.min():.4f}, Max={beta2_tensor.max():.4f}, Mean={beta2_tensor.mean():.4f}")
+                delattr(self, '_beta2_log')
+
         return loss

adv_optm-1.1.0.dev1/adv_optm/util/Kourkoutas.py

@@ -0,0 +1,108 @@
+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:
+            if not group.get('kourkoutas_beta', False):
+                continue
+            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)
+        self._layer_info_built = True
+
+    def prepare_step(self):
+        """
+        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()
+
+        if hasattr(self.optimizer, 'logging') and self.optimizer.logging:
+            if not hasattr(self.optimizer, '_beta2_log'):
+                self.optimizer._beta2_log = []
+
+        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]
+            
+            pooled_grad_norm = torch.sqrt(layer_state['sum_sq_accumulator'])
+            
+            r_ema = layer_state['r_ema_grad_norm']
+            r_ema.mul_(group['ema_alpha']).add_(pooled_grad_norm, alpha=1.0 - group['ema_alpha'])
+            
+            raw = pooled_grad_norm / (r_ema + group['tiny_spike'])
+            sun = raw / (1.0 + raw)
+            beta2_max = group['betas'][1]
+            beta2 = beta2_max - (beta2_max - group['beta2_min']) * sun
+            
+            layer_state['dynamic_beta2'] = beta2.item()
+            layer_state['sum_sq_accumulator'].zero_()
+
+            if hasattr(self.optimizer, 'logging') and self.optimizer.logging and hasattr(self.optimizer, '_beta2_log'):
+                self.optimizer._beta2_log.append(beta2.item())
+
+    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()
+            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.
+        """
+        layer_key = self.optimizer.layer_key_fn(p)
+        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)
+            }
+        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.
+        """
+        beta2_default = group['betas'][1]
+        if current_step < group['k_warmup_steps']:
+            return 0.5 * (group['beta2_min'] + beta2_default)
+        
+        layer_key = self.optimizer.layer_key_fn(p)
+        return self.layer_state.get(layer_key, {}).get('dynamic_beta2', beta2_default)

{adv_optm-1.0.6 → adv_optm-1.1.0.dev1}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-1.0.6 → adv_optm-1.1.0.dev1}/adv_optm.egg-info/SOURCES.txt

@@ -16,6 +16,7 @@ adv_optm/optim/Simplified_AdEMAMix.py
 adv_optm/optim/__init__.py
 adv_optm/util/BF16_Stochastic_Rounding.py
 adv_optm/util/Effective_Shape.py
+adv_optm/util/Kourkoutas.py
 adv_optm/util/NNMF.py
 adv_optm/util/One_Bit_Boolean.py
 adv_optm/util/OrthoGrad.py

{adv_optm-1.0.6 → adv_optm-1.1.0.dev1}/setup.py

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