adv-optm 1.2.dev3__tar.gz → 1.2.dev5__tar.gz

{adv_optm-1.2.dev3 → adv_optm-1.2.dev5}/PKG-INFO

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

{adv_optm-1.2.dev3 → adv_optm-1.2.dev5}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "1.2.dev3"
+__version__ = "1.2.dev5"

{adv_optm-1.2.dev3 → adv_optm-1.2.dev5}/adv_optm/optim/AdaMuon_adv.py

@@ -135,7 +135,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             print("Warning: nesterov is incompatible with Simplified_AdEMAMix, Disabling cautious.")
             nesterov = False
 
-        defaults = {
+        muon_defaults = {
             "lr": lr, "betas": betas, "weight_decay": weight_decay,
             "eps": eps, "rms_target": rms_target, "ns_steps": ns_steps,
             "ns_eps": ns_eps, "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor,
@@ -154,22 +154,45 @@ class AdaMuon_adv(torch.optim.Optimizer):
         self.helper = None
         self.aux_adam = None
  
-        if self.MuonWithAuxAdam:
-            adam_kwargs = adam_kwargs or {}
-            # Create a delegate AdamW optimizer to get its default hyperparameters.
-            self.aux_adam = AdamW_adv(
-                [],
-                lr=muon_adam_lr,
-                **adam_kwargs,
-                _is_delegate=True
-            )
-            # Update the defaults dictionary
-            defaults.update(self.aux_adam.defaults)
-        
-        super().__init__(params, defaults)
+        if not self.MuonWithAuxAdam:
+            super().__init__(params, muon_defaults)
+            return
 
-        if self.MuonWithAuxAdam:
-            self.helper = MuonAdamHelper(self, layer_key_fn)
+        # HYBRID OPTIMIZER LOGIC
+        adam_kwargs = adam_kwargs or {}
+        self.aux_adam = AdamW_adv(
+            [],
+            lr=muon_adam_lr,
+            **adam_kwargs,
+            _is_delegate=True
+        )
+        adam_defaults = self.aux_adam.defaults
+
+        final_param_groups = []
+        _layer_key_fn = layer_key_fn if layer_key_fn is not None else lambda p: 'muon'
+
+        for group in params:
+            # All params in a group are of the same type
+            first_param = group['params'][0]
+            key = _layer_key_fn(first_param)
+            optim_type = 'adam' if key == 'adam' else 'muon'
+
+            new_group = group.copy()
+            defaults_to_use = adam_defaults if optim_type == 'adam' else muon_defaults
+
+            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, {})
+        
+        # Now that self is initialized, create the helper
+        self.helper = MuonAdamHelper(self, layer_key_fn)
         
 
     @property
@@ -196,21 +219,20 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
     @torch.no_grad()
     def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
+        if group['_kourkoutas_beta'] and self._kourkoutas_helper is None:
+            self._kourkoutas_helper = KourkoutasHelper(self)
+
         if self.MuonWithAuxAdam:
             optim_type = self.helper.get_optimizer_type(p)
             if optim_type == 'adam':
                 # Delegate to the AdamW_adv optimizer's logic.
                 # We need to temporarily "lend" our state and param_groups
-                # to the delegate so it has the full context to work with,
-                # especially for features like Kourkoutas-beta.
                 self.aux_adam.state = self.state
                 self.aux_adam.param_groups = self.param_groups
+
                 self.aux_adam.step_parameter(p, group, i)
                 return
 
-        if group['_kourkoutas_beta'] and self._kourkoutas_helper is None:
-            self._kourkoutas_helper = KourkoutasHelper(self)
-
         if p.grad is None:
             return
 

{adv_optm-1.2.dev3 → adv_optm-1.2.dev5}/adv_optm/optim/Muon_adv.py

@@ -100,7 +100,7 @@ class Muon_adv(torch.optim.Optimizer):
             print("Warning: nesterov is incompatible with Simplified_AdEMAMix, Disabling cautious.")
             nesterov = False
 
-        defaults = {
+        muon_defaults = {
             "lr": lr, "beta1": beta1, "weight_decay": weight_decay,
             "nesterov": nesterov, "ns_steps": ns_steps, "ns_eps": ns_eps,
             "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor,
@@ -114,23 +114,41 @@ class Muon_adv(torch.optim.Optimizer):
         self.helper = None
         self.aux_adam = None
  
-        if self.MuonWithAuxAdam:
-            adam_kwargs = adam_kwargs or {}
-            # Create a delegate AdamW optimizer to get its default hyperparameters.
-            self.aux_adam = AdamW_adv(
-                [],
-                lr=muon_adam_lr,
-                **adam_kwargs,
-                _is_delegate=True
-            )
-            # Update the defaults dictionary
-            defaults.update(self.aux_adam.defaults)
-        
-        super().__init__(params, defaults)
+        if not self.MuonWithAuxAdam:
+            super().__init__(params, muon_defaults)
+            return
 
-        if self.MuonWithAuxAdam:
-            self.helper = MuonAdamHelper(self, layer_key_fn)
+        # HYBRID OPTIMIZER LOGIC
+        adam_kwargs = adam_kwargs or {}
+        self.aux_adam = AdamW_adv(
+            [],
+            lr=muon_adam_lr,
+            **adam_kwargs,
+            _is_delegate=True
+        )
+        adam_defaults = self.aux_adam.defaults
+
+        final_param_groups = []
+        _layer_key_fn = layer_key_fn if layer_key_fn is not None else lambda p: 'muon'
+
+        for group in params:
+            first_param = group['params'][0]
+            key = _layer_key_fn(first_param)
+            optim_type = 'adam' if key == 'adam' else 'muon'
+
+            new_group = group.copy()
+            defaults_to_use = adam_defaults if optim_type == 'adam' else muon_defaults
+
+            for key, value in defaults_to_use.items():
+                new_group.setdefault(key, value)
+
+            final_param_groups.append(new_group)
+
+        super().__init__(final_param_groups, {})
         
+        # Now that self is initialized, create the helper
+        self.helper = MuonAdamHelper(self, layer_key_fn)
+
 
     @property
     def supports_fused_back_pass(self):
@@ -335,4 +353,4 @@ class Muon_adv(torch.optim.Optimizer):
             for i, p in enumerate(group['params']):
                 self.step_parameter(p, group, i)
 
-        return loss
+        return loss

{adv_optm-1.2.dev3 → adv_optm-1.2.dev5}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-1.2.dev3 → adv_optm-1.2.dev5}/setup.py

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