adv-optm 2.1.dev2__tar.gz → 2.1.dev3__tar.gz

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/PKG-INFO

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

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "2.1.dev2"
+__version__ = "2.1.dev3"

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/optim/AdaMuon_adv.py

@@ -9,7 +9,7 @@ from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 from ..util import Muon_AuxAdam
 
-A = torch.as_tensor(4 / math.pi)
+A = 4 / math.pi
 
 class AdaMuon_adv(torch.optim.Optimizer):
     """
@@ -396,7 +396,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
                     del denom, vt_buf
 
                 # RMS-aligned scaling
-                step_scale = lr * A if group['use_atan2'] else lr
+                step_scale = lr * A if group['use_atan2'] and not group['normuon_variant'] else lr
                 rms_adjustment(update, group['rms_rescaling'], step_scale)
 
                 update = update.reshape(p.shape)
@@ -454,14 +454,18 @@ class AdaMuon_adv(torch.optim.Optimizer):
                     del denom
 
                 # RMS-aligned rescaling
-                step_scale = lr * A if group['use_atan2'] else lr
+                step_scale = lr * A if group['use_atan2'] and not group['normuon_variant'] else lr
                 rms_adjustment(update, group['rms_rescaling'], step_scale)
 
                 update = update.reshape(original_shape)
 
             param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
 
-        compiled_muon_step_parameter(state, grad, group, group['lr'], random_int_tensor)
+        if group.get('compiled_optimizer', False):
+            lr = torch.as_tensor(group['lr'])
+        else:
+            lr = group['lr']
+        compiled_muon_step_parameter(state, grad, group, lr, random_int_tensor)
 
     @torch.no_grad()
     def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/optim/AdamW_adv.py

@@ -10,7 +10,7 @@ from ..util.update_util import _grams_update, _cautious_update
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 
-A = torch.as_tensor(4 / math.pi)
+A = 4 / math.pi
 
 class AdamW_adv(torch.optim.Optimizer):
     """
@@ -249,6 +249,7 @@ class AdamW_adv(torch.optim.Optimizer):
         random_int_tensor = None
 
         if group.get('compiled_optimizer', False):
+            step_size = torch.as_tensor(step_size)
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/optim/Adopt_adv.py

@@ -9,7 +9,7 @@ from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 from ..util.update_util import _grams_update, _cautious_update
 
-A = torch.as_tensor(4 / math.pi)
+A = 4 / math.pi
 
 class Adopt_adv(torch.optim.Optimizer):
     """
@@ -270,14 +270,15 @@ class Adopt_adv(torch.optim.Optimizer):
         random_int_tensor = None
 
         if group.get('compiled_optimizer', False):
+            lr = torch.as_tensor(group['lr'])
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             step_param_fn = self._compiled_step_parameter
         else:
+            lr = group['lr']
             step_param_fn = self._step_parameter
 
-        lr = group['lr']
 
         step_param_fn(p, grad, state, group, lr, beta1, beta2, random_int_tensor)
 

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/optim/Lion_Prodigy_adv.py

@@ -226,6 +226,7 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
             d = torch.as_tensor(group['d'])
+            dlr = torch.as_tensor(group['dlr'])
             step_param_fn = self._compiled_step_parameter
         else:
             d = group['d']

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/optim/Muon_adv.py

@@ -399,7 +399,11 @@ class Muon_adv(torch.optim.Optimizer):
 
             param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
 
-        compiled_muon_step_parameter(state, grad, group, group['lr'], random_int_tensor)
+        if group.get('compiled_optimizer', False):
+            lr = torch.as_tensor(group['lr'])
+        else:
+            lr = group['lr']
+        compiled_muon_step_parameter(state, grad, group, lr, random_int_tensor)
 
     @torch.no_grad()
     def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/optim/Prodigy_adv.py

@@ -11,7 +11,7 @@ from ..util.Kourkoutas import KourkoutasHelper
 from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state
 from ..util.update_util import _grams_update, _cautious_update
 
-A = torch.as_tensor(4 / math.pi)
+A = 4 / math.pi
 
 class Prodigy_adv(torch.optim.Optimizer):
     """
@@ -343,6 +343,7 @@ class Prodigy_adv(torch.optim.Optimizer):
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
             d = torch.as_tensor(group['d'])
+            dlr = torch.as_tensor(dlr)
             step_param_fn = self._compiled_step_parameter
         else:
             d = group['d']

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/optim/Simplified_AdEMAMix.py

@@ -244,7 +244,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
-            sqrt_den_num = torch.as_tensor(sqrt_den_num)
+            lr = torch.as_tensor(lr)
             step_param_fn = self._compiled_step_parameter
         else:
             step_param_fn = self._step_parameter
@@ -289,10 +289,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
             state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(vt, signed=False)
             del vt
 
-            if group['use_bias_correction']:
-                update.mul_(sqrt_den_num)
-
-            update = update.view(p.shape).mul_(lr)
+            update = update.view(p.shape).mul_(lr * sqrt_den_num)
 
         else:  # Standard optimizer logic for non-factored tensors
             exp_avg_sq = state['exp_avg_sq']
@@ -308,8 +305,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
             update.div_(denom)
             del denom
 
-            update_scaling = lr * sqrt_den_num if group['use_bias_correction'] else lr
-            update.mul_(update_scaling)
+            update.mul_(lr * sqrt_den_num)
 
         param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
 

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/util/Kourkoutas.py

@@ -182,7 +182,7 @@ class KourkoutasHelper:
         """
         layer_key = self.optimizer.layer_key_fn(p)
 
-        if layer_key in self.layer_info:
+        if layer_key in self.layer_info and layer_key in self.layer_state:
             # Accumulate for the *next* step's prepare_step call
             self.layer_state[layer_key]['sum_sq_accumulator'] += torch.sum(grad.detach().pow(2)).float()
 

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm/util/Muon_AuxAdam.py

@@ -7,7 +7,7 @@ from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state
 from ..util.update_util import _grams_update, _cautious_update
 
-A = torch.as_tensor(4 / math.pi)
+A = 4 / math.pi
 
 @torch.no_grad()
 def _init_auxadam_state(self, p, group):
@@ -56,7 +56,7 @@ def _adam_step_parameter(self, p, grad, state, group, is_compiled, random_int_te
 
     if self.kourkoutas_helper:
         # Prepare Kourkoutas-β once per optimizer step.
-        self.kourkoutas_helper.maybe_prepare_step(step)
+        self.kourkoutas_helper.maybe_prepare_step(step, p.device)
         # Get the dynamic beta2_adam calculated in prepare_step()
         beta2_adam = self.kourkoutas_helper.get_beta2(p, group)
 
@@ -73,6 +73,9 @@ def _adam_step_parameter(self, p, grad, state, group, is_compiled, random_int_te
 
     step_size = group['lr'] / bias_correction1
 
+    if group.get('compiled_optimizer', False):
+        step_size = torch.as_tensor(step_size)
+
     @torch.compile(fullgraph=True, disable= not is_compiled)
     def compiled_muon_step_parameter(state, grad, group, step_size, sqrt_bias_correction2, random_int_tensor):
         if grad.dtype != torch.float32 and state.get('factored', False):

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-2.1.dev2 → adv_optm-2.1.dev3}/setup.py

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