adv-optm 1.2.dev4__tar.gz → 1.2.dev6__tar.gz

{adv_optm-1.2.dev4 → adv_optm-1.2.dev6}/PKG-INFO

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

{adv_optm-1.2.dev4 → adv_optm-1.2.dev6}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "1.2.dev4"
+__version__ = "1.2.dev6"

{adv_optm-1.2.dev4 → adv_optm-1.2.dev6}/adv_optm/optim/AdaMuon_adv.py

@@ -182,7 +182,11 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
             for key, value in defaults_to_use.items():
                 new_group.setdefault(key, value)
-
+            if '_kourkoutas_beta' not in new_group:
+                 if optim_type == 'adam':
+                     new_group['_kourkoutas_beta'] = False
+                 else:
+                     new_group['_kourkoutas_beta'] = muon_defaults['_kourkoutas_beta']
             final_param_groups.append(new_group)
 
         super().__init__(final_param_groups, {})
@@ -225,10 +229,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 # We need to temporarily "lend" our state and param_groups
                 self.aux_adam.state = self.state
                 self.aux_adam.param_groups = self.param_groups
-                
-                # Ensure the aux optimizer uses the same Kourkoutas helper instance.
-                if self._kourkoutas_helper is not None:
-                    self.aux_adam.kourkoutas_helper = self._kourkoutas_helper
 
                 self.aux_adam.step_parameter(p, group, i)
                 return

{adv_optm-1.2.dev4 → adv_optm-1.2.dev6}/adv_optm/util/Kourkoutas.py

@@ -86,9 +86,17 @@ class KourkoutasHelper:
         # 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
 
+        # The optimizer that owns this helper holds the master defaults for K-b.
+        # This is crucial in hybrid optimizers where some param_groups might not
+        # have all K-b keys populated, preventing KeyErrors.
+        master_defaults = self.optimizer.defaults
+
         for layer_key, info in self.layer_info.items():
             params, group = info['params'], info['group_ref']
 
+            if not group.get('kourkoutas_beta', False):
+                continue
+
             first_param_in_layer = info['params'][0]
             param_state = self.optimizer.state[first_param_in_layer]
 
@@ -100,6 +108,15 @@ class KourkoutasHelper:
             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)
 
+            # Use group-specific K-b settings, falling back to the optimizer's master defaults.
+            # This makes the helper robust against param groups that enable kourkoutas_beta
+            # but are missing the other required hyperparameters.
+            ema_alpha = group.get('ema_alpha', master_defaults['ema_alpha'])
+            beta2_max = group.get('betas', master_defaults['betas'])[1]
+            beta2_min = group.get('beta2_min', master_defaults['beta2_min'])
+            tiny_spike = group.get('tiny_spike', master_defaults['tiny_spike'])
+            k_warmup_steps = group.get('k_warmup_steps', master_defaults['k_warmup_steps'])
+
             r_ema_tensor = param_state['kourkoutas_r_ema']
             accumulator = self.layer_state[layer_key]['sum_sq_accumulator']
             
@@ -107,17 +124,16 @@ class KourkoutasHelper:
             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'])
+            r_ema_tensor.mul_(ema_alpha).add_(pooled_grad_norm, alpha=1.0 - 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']:
+            if current_step < k_warmup_steps:
                 beta2 = beta2_max
             else:
-                raw = pooled_grad_norm / (r_ema_tensor + group['tiny_spike'])
+                raw = pooled_grad_norm / (r_ema_tensor + tiny_spike)
                 sun = raw / (1.0 + raw)
-                beta2 = beta2_max - (beta2_max - group['beta2_min']) * sun
+                beta2 = beta2_max - (beta2_max - 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

{adv_optm-1.2.dev4 → adv_optm-1.2.dev6}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-1.2.dev4 → adv_optm-1.2.dev6}/setup.py

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