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

adv_optm/__init__.py

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

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,41 @@ 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)
+
+            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 +215,24 @@ 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
+                
+                # 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
 
-        if group['_kourkoutas_beta'] and self._kourkoutas_helper is None:
-            self._kourkoutas_helper = KourkoutasHelper(self)
-
         if p.grad is None:
             return
 

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.dist-info → adv_optm-1.2.dev4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 1.2.dev3
+Version: 1.2.dev4
 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.dist-info → adv_optm-1.2.dev4.dist-info}/RECORD

@@ -1,10 +1,10 @@
-adv_optm/__init__.py,sha256=L18OZGOBUDD4MFER_tFQXnf8zpfvmtqqJsS6IqgAcME,379
-adv_optm/optim/AdaMuon_adv.py,sha256=pTf9sMwhtR2AO7ABTrdH7mM7nz4Pn3EgxBqOR0hx9AM,20246
+adv_optm/__init__.py,sha256=bB7_VywKpvZbcGCjtZoF8giQgcUgoziISBgIaEUpcAw,379
+adv_optm/optim/AdaMuon_adv.py,sha256=s5UkR2YJ_Z10SiBokT97eq4tCHc2D8BEOFDx5AOMryQ,20983
 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
 adv_optm/optim/Lion_adv.py,sha256=aGNAplZlyXYgVllYcV_s4bK8iC4fv6EizFoWIMNLdBc,8299
-adv_optm/optim/Muon_adv.py,sha256=K4Iub3Xu8tJNSc3VvvIOLgD322HZn_Dzu3DL-4BH9oA,14417
+adv_optm/optim/Muon_adv.py,sha256=vB-Eeh0IqYMd3lkQvIPEbH256bTyYO73OgIzn0N2VCk,14985
 adv_optm/optim/Prodigy_adv.py,sha256=bmwuO8GrJHH4NaEaqE-ffcR9wHhQ57457xoN-P6hyks,25909
 adv_optm/optim/Simplified_AdEMAMix.py,sha256=sY-vThMVgADRh0ar9WHkrM2n8UcgQLQC1YV1Wx8uFz4,12983
 adv_optm/optim/__init__.py,sha256=hpUWE6CKtt_rvMdgQVb3PtjhfZAvAxTq6hp8H8rIpBo,489
@@ -17,8 +17,8 @@ adv_optm/util/Newton_Schulz.py,sha256=wJ_sKRaGVIsOofQ737my4ng494qX_pfgOqlDDmYtnC
 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.dev3.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-1.2.dev3.dist-info/METADATA,sha256=twwbf6sSxWaS3YSt7NuDwSoWkP4H5zByk7Odo_CsTDY,14022
-adv_optm-1.2.dev3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-1.2.dev3.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-1.2.dev3.dist-info/RECORD,,
+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,,