adv-optm 1.2.dev1__py3-none-any.whl → 1.2.dev2__py3-none-any.whl

adv_optm/__init__.py

@@ -18,4 +18,4 @@ __all__ = [
     "Muon_adv",
 ]
 
-__version__ = "1.2.dev1"
+__version__ = "1.2.dev2"

adv_optm/optim/AdamW_adv.py

@@ -107,6 +107,7 @@ class AdamW_adv(torch.optim.Optimizer):
         k_logging: int = 0,
         layer_key_fn: Optional[Callable] = None,
         nnmf_factor: bool = False,
+        _is_delegate: bool = False,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -137,10 +138,11 @@ class AdamW_adv(torch.optim.Optimizer):
         self.factored = nnmf_factor
         self.kourkoutas_beta = kourkoutas_beta
         self.layer_key_fn = layer_key_fn
-        super().__init__(params, defaults)
-
-        if self.kourkoutas_beta:
-            self.kourkoutas_helper = KourkoutasHelper(self)
+        if not _is_delegate:
+            super().__init__(params, defaults)
+        else:
+            self.defaults = defaults
+        self.kourkoutas_helper = None
 
     @property
     def supports_fused_back_pass(self):
@@ -158,6 +160,8 @@ class AdamW_adv(torch.optim.Optimizer):
     def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
         if p.grad is None:
             return
+        if group.get('kourkoutas_beta', False) and self.kourkoutas_helper is None:
+            self.kourkoutas_helper = KourkoutasHelper(self)
 
         grad = p.grad
         if grad.dtype != torch.float32 and self.factored:

adv_optm/optim/Muon_adv.py

@@ -1,5 +1,8 @@
 import torch
-from typing import Optional
+from typing import Optional, Callable
+
+from .AdamW_adv import AdamW_adv
+from ..util.MuonAdam_helper import MuonAdamHelper
 
 from ..util.BF16_Stochastic_Rounding import add_stochastic_
 from ..util.Newton_Schulz import _newton_schulz_iteration
@@ -18,6 +21,10 @@ class Muon_adv(torch.optim.Optimizer):
     This implementation is designed for 2D parameters (e.g., linear layers) and
     can handle other-dimensional parameters (e.g., 1D bias, 4D convolutional layers) by
     flattening/reshaping them.
+    
+    This version can also operate in a hybrid mode, using an auxiliary AdamW
+    optimizer for specific parameters (e.g., biases, norms, embeddings) as
+    defined by a `layer_key_fn`.
 
     Args:
         params (iterable): iterable of parameters to optimize or dicts defining
@@ -39,6 +46,16 @@ class Muon_adv(torch.optim.Optimizer):
             matrices to apply low-rank compression (default: True).
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
+        MuonWithAuxAdam (bool): If True, enables the hybrid optimizer mode.
+            Parameters designated by `layer_key_fn` will be optimized with
+            AdamW_adv instead of Muon. (default: False)
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
+            and returns a key. If the key is 'adam', the parameter is handled by
+            the auxiliary AdamW optimizer. All other keys are handled by Muon.
+            Only used when `MuonWithAuxAdam` is True. (default: None)
+        adam_kwargs (Optional[dict]): A dictionary of keyword arguments to pass
+            to the auxiliary AdamW_adv optimizer. Only used when
+            `MuonWithAuxAdam` is True. (default: None)
     """
 
     def __init__(
@@ -55,6 +72,11 @@ class Muon_adv(torch.optim.Optimizer):
         vector_reshape_muon: bool = False,
         vector_reshape: bool = True,
         nnmf_factor: bool = False,
+        # hybrid optimizer mode
+        MuonWithAuxAdam: bool = False,
+        layer_key_fn: Optional[Callable] = None,
+        muon_adam_lr: float = 1e-4,
+        adam_kwargs: Optional[dict] = None,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -73,8 +95,29 @@ class Muon_adv(torch.optim.Optimizer):
             "vector_reshape_muon": vector_reshape_muon,
         }
         self.stochastic_rounding = stochastic_rounding
+        
+        self.MuonWithAuxAdam = MuonWithAuxAdam
+        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 self.MuonWithAuxAdam:
+            self.helper = MuonAdamHelper(self, layer_key_fn)
+        
+
     @property
     def supports_fused_back_pass(self):
         return True
@@ -89,6 +132,18 @@ class Muon_adv(torch.optim.Optimizer):
 
     @torch.no_grad()
     def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
+        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 p.grad is None:
             return
 
@@ -244,4 +299,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/util/MuonAdam_helper.py

@@ -0,0 +1,31 @@
+from torch.optim import Optimizer
+from typing import Callable, Optional
+
+class MuonAdamHelper:
+    """
+    A helper class for Muon_adv to decide whether to use Muon or a delegate
+    AdamW optimizer for a given parameter based on a keying function.
+    """
+    def __init__(self, optimizer: Optimizer, layer_key_fn: Optional[Callable]):
+        if not hasattr(optimizer, 'param_groups'):
+            raise TypeError("optimizer must be a valid torch.optim.Optimizer instance.")
+        self.optimizer = optimizer
+        
+        if layer_key_fn is None:
+            # If no function is provided, default all parameters to 'muon'.
+            self.layer_key_fn = lambda p: 'muon'
+        else:
+            self.layer_key_fn = layer_key_fn
+
+    def get_optimizer_type(self, p: "torch.Tensor") -> str:
+        """
+        Returns the designated optimizer type ('adam' or 'muon') for a parameter.
+        
+        The user-provided layer_key_fn should return 'adam' for parameters
+        to be handled by the auxiliary AdamW optimizer. Any other return
+        value is treated as 'muon'.
+        """
+        key = self.layer_key_fn(p)
+        if key == 'adam':
+            return 'adam'
+        return 'muon'

{adv_optm-1.2.dev1.dist-info → adv_optm-1.2.dev2.dist-info}/METADATA

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

{adv_optm-1.2.dev1.dist-info → adv_optm-1.2.dev2.dist-info}/RECORD

@@ -1,22 +1,23 @@
-adv_optm/__init__.py,sha256=fdpoxqapAZbMiPax_P4Zm9PkN--71G0yds0q0V9oAbo,341
-adv_optm/optim/AdamW_adv.py,sha256=7vWfPS2J54U9ZKFQiNJ_l86PvITb0MQ61Fy4Fzmf1d4,17479
+adv_optm/__init__.py,sha256=THWhNF8-PI71K9Au4xAkuDs96YcEagJ-yT5r_g2-yKw,341
+adv_optm/optim/AdamW_adv.py,sha256=Zym0beeu0ye5_PgpAjpzcYghdPYFWs3gQzDmuPZVR80,17690
 adv_optm/optim/Adopt_adv.py,sha256=NXbtPrGm3tZr06cApi5oEHZ2F1zwss3tRi15SGnrYPc,21426
 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=eXqPL6GIWutBJpP7Yb_qIk7pGAjwfTAloCFRDhkRoUU,9908
+adv_optm/optim/Muon_adv.py,sha256=9K5YR3odaGfDDZzasletHRlqxG8xN9IXj6oiqx1CaEI,12423
 adv_optm/optim/Prodigy_adv.py,sha256=0_XG5YnMQTv-zJysJHlJniSo5kGYdX3p3o1e33HLt78,25897
 adv_optm/optim/Simplified_AdEMAMix.py,sha256=nEIA3yM11nBooKzHudB5l3x4UdFRBYRwiKVUkGmO0K8,12971
 adv_optm/optim/__init__.py,sha256=3o2XJ4J-PUq3rJM2mBnmuHwbKNb4LuW-Ig_9aBC0ycc,431
 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=woyJfX7l4eieeg0pC5XrILBLvwECwbD3a6ou1K6qjKU,8706
+adv_optm/util/MuonAdam_helper.py,sha256=llPCc9MBFen_wodbY4G2E17tBZky8clDiJSZLHkMva8,1236
 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.dev1.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-1.2.dev1.dist-info/METADATA,sha256=ofbAQu0ldYk8udMEC0jLcI9Ex2a6M8iaXheDTo4Un3M,14022
-adv_optm-1.2.dev1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-1.2.dev1.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-1.2.dev1.dist-info/RECORD,,
+adv_optm-1.2.dev2.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-1.2.dev2.dist-info/METADATA,sha256=JTCPGBJUd4JR7DU26AhX8qSPzWrSVtEwv9Au7I3iEPY,14022
+adv_optm-1.2.dev2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+adv_optm-1.2.dev2.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-1.2.dev2.dist-info/RECORD,,