adv-optm 1.0.5__py3-none-any.whl → 1.1.0__py3-none-any.whl

adv_optm/optim/Prodigy_adv.py

@@ -3,15 +3,18 @@ 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):
     """
-    Implements a factored Prodigy/AdamW algorithm.
+    Implements an advanced Prodigy algorithm.
     This is an advanced version of Prodigy with optional features like
     low-rank factorization of optimizer states (SMMF), OrthoGrad, etc.
 
@@ -85,6 +88,31 @@ 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).
+        d_limiter (bool): whether to clamp the new step size estimate (`d_hat`)
+            to prevent sudden, volatile increases in the adaptive step size (`d`).
+            (default: False)
+        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 beta2 value 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 +144,15 @@ class Prodigy_adv(torch.optim.Optimizer):
         fsdp_in_use: bool = False,
         slice_p: int = 11,
         prodigy_steps: int = 0,
+        d_limiter: bool = False,
+        # K-b parameters
+        kourkoutas_beta: bool = False,
+        beta2_min: float = 0.9,
+        ema_alpha: float = 0.95,
+        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,8 +178,10 @@ 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 Simplified_AdEMAMix and alpha_grad > 0:
-            # scales d_coef by alpha_grad, this force prodigy to behave well with Simplified_AdEMAMix
+        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 and not d_limiter:
+            # scales d_coef by alpha_grad, this force prodigy to behave well with Simplified_AdEMAMix.
             d_coef = d_coef/alpha_grad
 
         defaults = {
@@ -152,8 +191,10 @@ class Prodigy_adv(torch.optim.Optimizer):
             "beta3_ema": beta3_ema, "alpha": alpha, "t_alpha": t_alpha,
             "beta3": beta3, "d": d0, "d0": d0, "d_max": d0, "d_numerator": 0.0, "d_coef": d_coef,
             "growth_rate": growth_rate, "safeguard_warmup": safeguard_warmup, "k": 0, "slice_p": slice_p,
-            "fsdp_in_use": fsdp_in_use, "prodigy_steps": prodigy_steps,
-            "alpha_grad": alpha_grad, 
+            "fsdp_in_use": fsdp_in_use, "prodigy_steps": prodigy_steps, "d_limiter": d_limiter,
+            "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 +203,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 +227,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)
-        
-        k = g_group['k']
+            self.beta3 = math.sqrt(self.beta2_default)
+
         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()
@@ -211,7 +256,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         state = self.state[p]
 
         # State Initialization
-        if len(state) == 0:
+        if 'step' not in state:
             state['step'] = 0
 
             should_factor = (
@@ -258,14 +303,27 @@ 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)
+            beta3 = math.sqrt(beta2)
+        else:
+            beta2 = self.beta2_default
+            beta3 = self.beta3
+
         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 +353,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']))
@@ -308,11 +366,11 @@ class Prodigy_adv(torch.optim.Optimizer):
                 if self.beta1 > 0:
                     update = torch.add(mt, mt_slow, alpha=alpha_t)
                 else:
-                    update = torch.add(grad_reshaped, mt_slow, alpha=alpha_t)
+                    update = torch.add(grad_reshaped.mul(self.d), mt_slow, alpha=alpha_t)
             elif self.Simplified_AdEMAMix:
                 update = torch.add(mt, grad_reshaped, alpha=alpha_grad * self.d)
             else:
-                update = mt.clone() if self.beta1 > 0 else grad_reshaped.clone()
+                update = mt.clone() if self.beta1 > 0 else grad_reshaped.mul(self.d)
             del grad_reshaped
 
             if group['use_atan2']:
@@ -362,13 +420,13 @@ class Prodigy_adv(torch.optim.Optimizer):
                 if self.beta1 > 0:
                     update = torch.add(exp_avg, exp_avg_slow, alpha=alpha_t)
                 else:
-                    update = torch.add(grad, exp_avg_slow, alpha=alpha_t)
+                    update = torch.add(grad.mul(self.d), exp_avg_slow, alpha=alpha_t)
             elif self.Simplified_AdEMAMix:
                 update = torch.add(exp_avg, grad, alpha=alpha_grad * self.d)
             else:
-                update = exp_avg.clone() if self.beta1 > 0 else grad.clone()
+                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
@@ -393,7 +451,7 @@ class Prodigy_adv(torch.optim.Optimizer):
             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)
+            s.mul_(beta3).add_(grad_flat[::slice_p], alpha=alpha)
             self.d_denom += s.abs().sum().item()
 
             del s, p0, grad_flat, p_flat, alpha
@@ -431,7 +489,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
@@ -460,6 +517,8 @@ class Prodigy_adv(torch.optim.Optimizer):
             d_hat = self.d
             if global_d_denom > 0:
                 d_hat = d_coef * global_d_numerator / global_d_denom
+                if g_group['d_limiter']:
+                    d_hat = min(self.d * (2 ** 0.25), d_hat)
                 if self.d == g_group['d0']:
                     self.d = max(self.d, d_hat)
                 d_max = max(d_max, d_hat)

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 beta2 value 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.9,
+        ema_alpha: float = 0.95,
+        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
@@ -113,7 +152,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
         state = self.state[p]
 
         # State Initialization
-        if len(state) == 0:
+        if 'step' not in state:
             state['step'] = 0
 
             should_factor = (
@@ -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,171 @@
+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
+        self.layer_state = {}
+
+        self.layer_info = {}
+        self._layer_info_built = False
+        self._current_step_prepared = -1
+
+        # Store stats for external logging (e.g., TensorBoard)
+        self.last_beta2_stats = {}
+
+        # This ensures the map is complete before the first backward pass,
+        # making it compatible with fused back pass mechanisms.
+        self._build_layer_info_if_needed()
+
+        if self.optimizer.param_groups[0].get('k_logging', 0) > 0:
+            self.print_layer_info()
+
+    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']:
+                # The mapping is static and should not depend on the presence of a gradient.
+                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 print_layer_info(self):
+        """Prints the contents of self.layer_info for debugging."""
+        print("\n--- BEGIN self.layer_info DUMP ---")
+        if not self.layer_info:
+            print("Layer info is empty. Make sure the optimizer has parameters.")
+            return
+
+        for layer_key, info in self.layer_info.items():
+            param_count = len(info['params'])
+            first_param_details = ""
+            if param_count > 0:
+                p = info['params'][0]
+                first_param_details = f" (Example param shape: {list(p.shape)}, dtype: {p.dtype})"
+            
+            print(f"Key: {layer_key}, Params: {param_count}{first_param_details}")
+
+        print("--- END self.layer_info DUMP ---\n")
+
+    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.
+        """
+        
+        beta2_log = []
+        first_layer_key = next(iter(self.layer_info), None)
+        # These are just for the sample log, initialize them
+        sun, pooled_grad_norm, prev_r_ema_val, r_ema_tensor = (torch.tensor(0.0),)*4
+
+        for layer_key, info in self.layer_info.items():
+            params, group = info['params'], info['group_ref']
+
+            first_param_in_layer = info['params'][0]
+            param_state = self.optimizer.state[first_param_in_layer]
+
+            if layer_key not in self.layer_state:
+                self.layer_state[layer_key] = {
+                    'sum_sq_accumulator': torch.tensor(0.0, device=first_param_in_layer.device, dtype=torch.float32)
+                }
+            
+            if 'kourkoutas_r_ema' not in param_state:
+                param_state['kourkoutas_r_ema'] = torch.tensor(0.0, device=first_param_in_layer.device, dtype=torch.float32)
+
+            r_ema_tensor = param_state['kourkoutas_r_ema']
+            accumulator = self.layer_state[layer_key]['sum_sq_accumulator']
+            
+            pooled_grad_norm = torch.sqrt(accumulator)
+            prev_r_ema_val = r_ema_tensor.item() # for logging
+            
+            # Update the persistent EMA tensor in-place.
+            r_ema_tensor.mul_(group['ema_alpha']).add_(pooled_grad_norm, alpha=1.0 - group['ema_alpha'])
+            
+            beta2_max = group['betas'][1]
+            sun = torch.tensor(0.0, device=r_ema_tensor.device) # Default sun to 0 for warmup
+            
+            if current_step < group['k_warmup_steps']:
+                beta2 = beta2_max
+            else:
+                raw = pooled_grad_norm / (r_ema_tensor + group['tiny_spike'])
+                sun = raw / (1.0 + raw)
+                beta2 = beta2_max - (beta2_max - group['beta2_min']) * sun
+
+            # Store the final calculated beta2 in the helper's transient state for this step.
+            self.layer_state[layer_key]['dynamic_beta2'] = beta2.item() if isinstance(beta2, torch.Tensor) else beta2
+            
+            # Reset the accumulator for the next optimizer step.
+            accumulator.zero_()
+
+            beta2_log.append(self.layer_state[layer_key]['dynamic_beta2'])
+
+        # Always compute stats for TensorBoard
+        if beta2_log:
+            beta2_tensor = torch.tensor(beta2_log, device='cpu')
+            self.last_beta2_stats = {
+                'min': beta2_tensor.min().item(),
+                'max': beta2_tensor.max().item(),
+                'mean': beta2_tensor.mean().item(),
+            }
+
+        # Handle periodic console logging
+        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
+        if is_logging_step and self.last_beta2_stats:
+            if first_layer_key:
+                print(f"\n[Kourkoutas-β Debug] Step {current_step + 1} - Sample Layer '{first_layer_key}':")
+                print(f"  - Grad Norm: {pooled_grad_norm.item():.4e}, Prev EMA: {prev_r_ema_val:.4e}, New EMA: {r_ema_tensor.item():.4e}")
+                print(f"  - Sunspike: {sun.item():.4f}, Dynamic Beta2: {self.layer_state[first_layer_key]['dynamic_beta2']:.4f}")
+            print(f"[Kourkoutas-β Debug] Step {current_step + 1} Overall Beta2 Stats: Min={self.last_beta2_stats['min']:.4f}, Max={self.last_beta2_stats['max']:.4f}, Mean={self.last_beta2_stats['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.
+        """
+        layer_key = self.optimizer.layer_key_fn(p)
+
+        if layer_key in self.layer_info:
+            # Initialize the transient state for this layer if it's the first time in the step.
+            if layer_key not in self.layer_state:
+                    self.layer_state[layer_key] = {
+                    '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.5.dist-info → adv_optm-1.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 1.0.5
+Version: 1.1.0
 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.dist-info/RECORD

@@ -0,0 +1,20 @@
+adv_optm/__init__.py,sha256=lNp6_DdCvw-0zok2UdMkaEyVLZIlMRSKgBp-hJ15Hao,306
+adv_optm/optim/AdamW_adv.py,sha256=ddEUVOif1gfZPgEJNrEGZ2wnha4MPMWw5ppPd8acQ3o,17457
+adv_optm/optim/Adopt_adv.py,sha256=fhH3hS9K6z5Blxc7NFfzpCrUGbl9EQnwLPmKDxBC1zg,21415
+adv_optm/optim/Lion_Prodigy_adv.py,sha256=aJ9orEEw0QYbrDzn1be0SHvOBlIkLwWG9RpWFuNMskM,13163
+adv_optm/optim/Lion_adv.py,sha256=aGNAplZlyXYgVllYcV_s4bK8iC4fv6EizFoWIMNLdBc,8299
+adv_optm/optim/Prodigy_adv.py,sha256=4O7BLGhqLW46Ff3UN9JfrktHonCYDy3ojHUfW8jtaDs,25940
+adv_optm/optim/Simplified_AdEMAMix.py,sha256=gPjMhKulzmAeO42foe-d7xW0AcB50vKFYsvHgxbD3uc,12949
+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=DCsIcZ1sEeSwthN8KZH7OTKoIZJ3ah4t5DNiqxsSuCk,8344
+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.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-1.1.0.dist-info/METADATA,sha256=dwRwKQykba-7TP6a94qpOg6xz450QESAi5E8AnEV-iM,8422
+adv_optm-1.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+adv_optm-1.1.0.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-1.1.0.dist-info/RECORD,,

adv_optm-1.0.5.dist-info/RECORD

@@ -1,19 +0,0 @@
-adv_optm/__init__.py,sha256=9sM4fP1pj9divFhLVUzHbBWe50H82H3FYIGVIHTHpkg,306
-adv_optm/optim/AdamW_adv.py,sha256=aTuYcJgd_EcZOrs6TDgBrBKw3wtU5LPzE5WvTBDDeEo,14317
-adv_optm/optim/Adopt_adv.py,sha256=KdEVSl2w1gRXFtz2fwCVT4i9inTspp-PQq3mobpa-9A,17476
-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=8XUpu19BaBmHb-R9K3jgwySDbtVaLU1_Drtttc_zITs,22461
-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.5.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-1.0.5.dist-info/METADATA,sha256=ouxI4cwBQ2IPuOjrkA478XwSetGP6ku51vW1QxHIGcY,8422
-adv_optm-1.0.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-1.0.5.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-1.0.5.dist-info/RECORD,,