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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 1.1.0.dev1
+Version: 1.1.0.dev3
 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.dev1 → adv_optm-1.1.0.dev3}/adv_optm/__init__.py

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

{adv_optm-1.1.0.dev1 → adv_optm-1.1.0.dev3}/adv_optm/optim/AdamW_adv.py

@@ -128,18 +128,17 @@ class AdamW_adv(torch.optim.Optimizer):
             "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,
+            "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)
 
-        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)
 
@@ -207,13 +206,15 @@ class AdamW_adv(torch.optim.Optimizer):
                     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)
-        
-        beta1, beta2 = group['betas']
-        if group['kourkoutas_beta']:
+            # Get the dynamic beta2 calculated in prepare_step()
             beta2 = self.kourkoutas_helper.get_beta2(p, group, current_step)
 
         step = state['step'] + 1
@@ -366,14 +367,4 @@ class AdamW_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.1.0.dev1 → adv_optm-1.1.0.dev3}/adv_optm/optim/Adopt_adv.py

@@ -157,7 +157,7 @@ class Adopt_adv(torch.optim.Optimizer):
             "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,
+            "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
         }
         self.clip_lambda = clip_lambda
         self.stochastic_rounding = stochastic_rounding
@@ -168,11 +168,10 @@ 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)
 
-        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)
 
@@ -238,13 +237,15 @@ class Adopt_adv(torch.optim.Optimizer):
                     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)
-
-        beta1, beta2 = group['betas']
-        if group['kourkoutas_beta']:
+            # 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).
@@ -257,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"]
 
@@ -436,10 +437,4 @@ class Adopt_adv(torch.optim.Optimizer):
             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-1.1.0.dev3}/adv_optm/optim/Prodigy_adv.py

@@ -189,7 +189,7 @@ class Prodigy_adv(torch.optim.Optimizer):
             "fsdp_in_use": fsdp_in_use, "prodigy_steps": prodigy_steps,
             "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,
+            "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
@@ -198,14 +198,13 @@ class Prodigy_adv(torch.optim.Optimizer):
         self.Simplified_AdEMAMix = Simplified_AdEMAMix
         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
+        self.layer_key_fn = layer_key_fn
+
+        super().__init__(params, defaults)
         if self.kourkoutas_beta:
             self.kourkoutas_helper = KourkoutasHelper(self)
-
         self.init_step()
 
     @property
@@ -229,7 +228,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         self.beta3 = g_group['beta3']
         if self.beta3 is None:
             self.beta3 = math.sqrt(self.beta2_default)
-        
+
         self.d = g_group['d']
         lr = g_group['lr']
 
@@ -301,21 +300,25 @@ class Prodigy_adv(torch.optim.Optimizer):
 
         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)
-
-        beta2 = self.beta2_default
-        if group['kourkoutas_beta']:
+            # 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"]
 
@@ -443,7 +446,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
@@ -481,16 +484,6 @@ 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()
         return loss

{adv_optm-1.1.0.dev1 → adv_optm-1.1.0.dev3}/adv_optm/optim/Simplified_AdEMAMix.py

@@ -116,15 +116,14 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
             "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,
+            "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)
 
-        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)
 
@@ -189,17 +188,19 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
                 state['num_sum'] = 1.0
                 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_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
@@ -294,14 +295,4 @@ 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-1.1.0.dev3}/adv_optm/util/Kourkoutas.py

@@ -18,6 +18,10 @@ class KourkoutasHelper:
         self._layer_info_built = False
         self._current_step_prepared = -1
 
+        # 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()
+
     def _build_layer_info_if_needed(self):
         """Builds a map of layers and the parameters they contain."""
         if self._layer_info_built:
@@ -28,26 +32,31 @@ class KourkoutasHelper:
             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
+                # 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 prepare_step(self):
+    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()
 
-        if hasattr(self.optimizer, 'logging') and self.optimizer.logging:
-            if not hasattr(self.optimizer, '_beta2_log'):
-                self.optimizer._beta2_log = []
+        # 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']
@@ -60,28 +69,47 @@ class KourkoutasHelper:
             
             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'])
             
-            raw = pooled_grad_norm / (r_ema + group['tiny_spike'])
-            sun = raw / (1.0 + raw)
+            sun = torch.tensor(0.0, device=r_ema.device) # Default sun to 0 for warmup
             beta2_max = group['betas'][1]
-            beta2 = beta2_max - (beta2_max - group['beta2_min']) * sun
-            
-            layer_state['dynamic_beta2'] = beta2.item()
+
+            # --- 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 hasattr(self.optimizer, 'logging') and self.optimizer.logging and hasattr(self.optimizer, '_beta2_log'):
-                self.optimizer._beta2_log.append(beta2.item())
+            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()
+            self.prepare_step(current_step)
             self._current_step_prepared = current_step
 
     def accumulate_gradient_sq_norm(self, p: torch.Tensor, grad: torch.Tensor):
@@ -89,20 +117,20 @@ class KourkoutasHelper:
         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()
+
+        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.
         """
-        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)
+        # 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.1.0.dev1 → adv_optm-1.1.0.dev3}/adv_optm.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 1.1.0.dev1
+Version: 1.1.0.dev3
 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.dev1 → adv_optm-1.1.0.dev3}/setup.py

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