adv-optm 1.2.dev1__tar.gz → 1.2.dev2__tar.gz

{adv_optm-1.2.dev1 → adv_optm-1.2.dev2}/PKG-INFO

@@ -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 → adv_optm-1.2.dev2}/adv_optm/__init__.py

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

{adv_optm-1.2.dev1 → adv_optm-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-1.2.dev1 → adv_optm-1.2.dev2}/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-1.2.dev2/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 → adv_optm-1.2.dev2}/adv_optm.egg-info/PKG-INFO

@@ -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 → adv_optm-1.2.dev2}/adv_optm.egg-info/SOURCES.txt

@@ -18,6 +18,7 @@ adv_optm/optim/__init__.py
 adv_optm/util/BF16_Stochastic_Rounding.py
 adv_optm/util/Effective_Shape.py
 adv_optm/util/Kourkoutas.py
+adv_optm/util/MuonAdam_helper.py
 adv_optm/util/NNMF.py
 adv_optm/util/Newton_Schulz.py
 adv_optm/util/One_Bit_Boolean.py

{adv_optm-1.2.dev1 → adv_optm-1.2.dev2}/setup.py

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