adv-optm 1.2.dev4__py3-none-any.whl → 1.2.dev6__py3-none-any.whl

adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "1.2.dev4"
+__version__ = "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/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.dist-info → adv_optm-1.2.dev6.dist-info}/METADATA

@@ -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.dist-info → adv_optm-1.2.dev6.dist-info}/RECORD

@@ -1,5 +1,5 @@
-adv_optm/__init__.py,sha256=bB7_VywKpvZbcGCjtZoF8giQgcUgoziISBgIaEUpcAw,379
-adv_optm/optim/AdaMuon_adv.py,sha256=s5UkR2YJ_Z10SiBokT97eq4tCHc2D8BEOFDx5AOMryQ,20983
+adv_optm/__init__.py,sha256=PXTst9vLSSIgeIiqqsDcq1FehR-HFUCBT7cP6bDTteA,379
+adv_optm/optim/AdaMuon_adv.py,sha256=hTGSH8wzmQ-NYIcqV6EAEbqCxxfEwmmMWaIadX1qiuQ,21009
 adv_optm/optim/AdamW_adv.py,sha256=7IvdD1rqYeHZwQCZU9X0H7x87MCKcHQ5M68GLuMCkvE,17702
 adv_optm/optim/Adopt_adv.py,sha256=C2FsEZGvCk9q4YNKAj0qIxdZ5AfPlda-1lIpSX0a1nE,21256
 adv_optm/optim/Lion_Prodigy_adv.py,sha256=LEA3UYJpPeFnmxeniLNv1u2LKKj4ufx3Bq_MLw-nWXk,14617
@@ -10,15 +10,15 @@ adv_optm/optim/Simplified_AdEMAMix.py,sha256=sY-vThMVgADRh0ar9WHkrM2n8UcgQLQC1YV
 adv_optm/optim/__init__.py,sha256=hpUWE6CKtt_rvMdgQVb3PtjhfZAvAxTq6hp8H8rIpBo,489
 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=MDQaNVH8jqzaefks2RShveo44dpYDz88WStwUJ3iF0s,8724
+adv_optm/util/Kourkoutas.py,sha256=WPAjxaH9pGVtLK_QJcwjkJOnN02Hfyu0F2T90hbhtqo,9662
 adv_optm/util/MuonAdam_helper.py,sha256=7rnNMujZVDaqo1g22QscMyPlZvIHQQSLHMED9_I8QWU,1250
 adv_optm/util/NNMF.py,sha256=yRf5IP5Sjq0Uf0DxN0Q8NxEGSdD-f1ULziLVDOjY8K4,639
 adv_optm/util/Newton_Schulz.py,sha256=wJ_sKRaGVIsOofQ737my4ng494qX_pfgOqlDDmYtnCg,1377
 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=jAaUfaAjFrTJ6-Q915ezAbq0efRbpYjriW2OdeCbSzo,433
-adv_optm-1.2.dev4.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-1.2.dev4.dist-info/METADATA,sha256=jNczVxIPq0LuusXuGrZ23CQ4CrMNOfJdBDpDQgulMUw,14022
-adv_optm-1.2.dev4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-1.2.dev4.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-1.2.dev4.dist-info/RECORD,,
+adv_optm-1.2.dev6.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-1.2.dev6.dist-info/METADATA,sha256=IIuYXoR2_uWQo7LjFO-MnBymyygc-Ntd9Fh0HHxnieA,14022
+adv_optm-1.2.dev6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+adv_optm-1.2.dev6.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-1.2.dev6.dist-info/RECORD,,