adv-optm 1.2.dev8__tar.gz → 1.2.dev10__tar.gz

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/PKG-INFO

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

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "1.2.dev8"
+__version__ = "1.2.dev10"

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm/optim/AdaMuon_adv.py

@@ -1,8 +1,7 @@
 import torch
-from typing import Optional, Callable
+from typing import Optional
 
 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
@@ -73,10 +72,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
         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)
@@ -106,7 +101,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
         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,
     ):
@@ -132,10 +126,10 @@ class AdaMuon_adv(torch.optim.Optimizer):
             "low_rank_ortho": low_rank_ortho, "ortho_rank": ortho_rank,
         }
         self.stochastic_rounding = stochastic_rounding
+
         self.MuonWithAuxAdam = MuonWithAuxAdam
-        self.helper = None
         self.aux_adam = None
- 
+
         if not self.MuonWithAuxAdam:
             super().__init__(params, muon_defaults)
             return
@@ -151,26 +145,17 @@ class AdaMuon_adv(torch.optim.Optimizer):
         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()
+            optim_type = group.get('optim_type', 'muon')
             defaults_to_use = adam_defaults if optim_type == 'adam' else muon_defaults
 
+            new_group = group.copy()
             for key, value in defaults_to_use.items():
                 new_group.setdefault(key, value)
             final_param_groups.append(new_group)
 
         super().__init__(final_param_groups, muon_defaults)
         
-        # Now that self is initialized, create the helper
-        self.helper = MuonAdamHelper(self, layer_key_fn)
-        
 
     @property
     def supports_fused_back_pass(self):
@@ -197,13 +182,14 @@ class AdaMuon_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)
+            optim_type = group.get('optim_type')
             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
 
@@ -327,7 +313,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
             # RMS-aligned rescaling
             rms_target = group['rms_target']
             num_elements = update.numel()
-            scaling_factor = rms_target * (num_elements ** 0.5) / (update.norm())
+            # Add eps to prevent division by zero
+            scaling_factor = rms_target * (num_elements ** 0.5) / (update.norm() + group['eps'])
 
             update.mul_(scaling_factor)
             update = update.view(p.shape).mul_(group['lr'])
@@ -422,7 +409,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 # RMS-aligned rescaling
                 rms_target = group['rms_target']
                 num_elements = update.numel()
-                scaling_factor = rms_target * (num_elements ** 0.5) / (update.norm())
+                # Add eps to prevent division by zero
+                scaling_factor = rms_target * (num_elements ** 0.5) / (update.norm() + group['eps'])
 
                 update.mul_(scaling_factor)
                 del num_elements, scaling_factor

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm/optim/AdamW_adv.py

@@ -73,7 +73,7 @@ class AdamW_adv(torch.optim.Optimizer):
             logging of Kourkoutas-β statistics (min, max, mean of `β₂` across layers)
             every logging steps. Useful for debugging and tuning. Set to 0 to disable
             logging (default: 0). 
-        layer_key_kb_fn (Optional[Callable]): A function that takes a parameter `p`
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
             and returns a unique, hashable key representing its "layer" or "bucket".
             If `None`, parameters are bucketed by their memory ID (tensor-wise).
             (default: None)
@@ -105,7 +105,7 @@ class AdamW_adv(torch.optim.Optimizer):
         tiny_spike: float = 1e-9,
         k_warmup_steps: int = 0,
         k_logging: int = 0,
-        layer_key_kb_fn: Optional[Callable] = None,
+        layer_key_fn: Optional[Callable] = None,
         nnmf_factor: bool = False,
         _is_delegate: bool = False,
     ):
@@ -137,7 +137,7 @@ class AdamW_adv(torch.optim.Optimizer):
         self.use_AdEMAMix = use_AdEMAMix
         self.factored = nnmf_factor
         self.kourkoutas_beta = kourkoutas_beta
-        self.layer_key_kb_fn = layer_key_kb_fn
+        self.layer_key_fn = layer_key_fn
         if not _is_delegate:
             super().__init__(params, defaults)
         else:
@@ -244,6 +244,7 @@ class AdamW_adv(torch.optim.Optimizer):
 
         if state['factored']:
             d1, d2 = state['effective_shape']
+            grad_reshaped = grad.view(d1, d2)
 
             # Reconstruct momentum from previous step's factors
             if beta1 > 0:
@@ -253,7 +254,6 @@ class AdamW_adv(torch.optim.Optimizer):
                     torch.where(unpacked_sign, mt, -mt, out=mt)
                     del unpacked_sign
                 # Update momentum in full-size
-                grad_reshaped = grad.view(d1, d2)
                 mt.mul_(beta1).add_(grad_reshaped, alpha=1.0 - beta1)
                 if self.grams_moment:
                     mt.copy_(grad_reshaped.sign() * mt.abs())

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm/optim/Adopt_adv.py

@@ -91,7 +91,7 @@ class Adopt_adv(torch.optim.Optimizer):
             logging of Kourkoutas-β statistics (min, max, mean of `β₂` across layers)
             every logging steps. Useful for debugging and tuning. Set to 0 to disable
             logging (default: 0). 
-        layer_key_kb_fn (Optional[Callable]): A function that takes a parameter `p`
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
             and returns a unique, hashable key representing its "layer" or "bucket".
             If `None`, parameters are bucketed by their memory ID (tensor-wise).
             (default: None)
@@ -125,7 +125,7 @@ class Adopt_adv(torch.optim.Optimizer):
         tiny_spike: float = 1e-9,
         k_warmup_steps: int = 0,
         k_logging: int = 0,
-        layer_key_kb_fn: Optional[Callable] = None,
+        layer_key_fn: Optional[Callable] = None,
         nnmf_factor: bool = False,
     ):
         if not (lr >= 0.0):
@@ -166,7 +166,7 @@ class Adopt_adv(torch.optim.Optimizer):
         self.Simplified_AdEMAMix = Simplified_AdEMAMix
         self.factored = nnmf_factor
         self.kourkoutas_beta = kourkoutas_beta
-        self.layer_key_kb_fn = layer_key_kb_fn
+        self.layer_key_fn = layer_key_fn
         super().__init__(params, defaults)
 
         if self.kourkoutas_beta:

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm/optim/Muon_adv.py

@@ -1,8 +1,6 @@
 import torch
-from typing import Optional, Callable
-
+from typing import Optional
 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
@@ -74,10 +72,6 @@ class Muon_adv(torch.optim.Optimizer):
         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)
@@ -110,7 +104,6 @@ class Muon_adv(torch.optim.Optimizer):
         normuon_atan2: 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,
     ):
@@ -145,9 +138,8 @@ class Muon_adv(torch.optim.Optimizer):
         self.stochastic_rounding = stochastic_rounding
 
         self.MuonWithAuxAdam = MuonWithAuxAdam
-        self.helper = None
         self.aux_adam = None
- 
+
         if not self.MuonWithAuxAdam:
             super().__init__(params, muon_defaults)
             return
@@ -163,25 +155,16 @@ class Muon_adv(torch.optim.Optimizer):
         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()
+            optim_type = group.get('optim_type', 'muon')
             defaults_to_use = adam_defaults if optim_type == 'adam' else muon_defaults
 
+            new_group = group.copy()
             for key, value in defaults_to_use.items():
                 new_group.setdefault(key, value)
-
             final_param_groups.append(new_group)
 
         super().__init__(final_param_groups, muon_defaults)
-        
-        # Now that self is initialized, create the helper
-        self.helper = MuonAdamHelper(self, layer_key_fn)
 
 
     @property
@@ -209,7 +192,7 @@ 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)
+            optim_type = group.get('optim_type')
             if optim_type == 'adam':
                 # Delegate to the AdamW_adv optimizer's logic.
                 # We need to temporarily "lend" our state and param_groups

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm/optim/Prodigy_adv.py

@@ -109,7 +109,7 @@ class Prodigy_adv(torch.optim.Optimizer):
             logging of Kourkoutas-β statistics (min, max, mean of `β₂` across layers)
             every logging steps. Useful for debugging and tuning. Set to 0 to disable
             logging (default: 0). 
-        layer_key_kb_fn (Optional[Callable]): A function that takes a parameter `p`
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
             and returns a unique, hashable key representing its "layer" or "bucket".
             If `None`, parameters are bucketed by their memory ID (tensor-wise).
             (default: None)
@@ -152,7 +152,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         tiny_spike: float = 1e-9,
         k_warmup_steps: int = 0,
         k_logging: int = 0,
-        layer_key_kb_fn: Optional[Callable] = None,
+        layer_key_fn: Optional[Callable] = None,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -205,7 +205,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         self.fsdp_in_use = fsdp_in_use
         
         self.kourkoutas_beta = kourkoutas_beta
-        self.layer_key_kb_fn = layer_key_kb_fn
+        self.layer_key_fn = layer_key_fn
 
         super().__init__(params, defaults)
         if self.kourkoutas_beta:
@@ -516,7 +516,10 @@ class Prodigy_adv(torch.optim.Optimizer):
             if global_d_denom > 0:
                 d_hat = d_coef * global_d_numerator / global_d_denom
                 if g_group.get('d_limiter', False):
-                    d_hat = min(self.d * (2 ** 0.25), d_hat)
+                    if g_group.get('Simplified_AdEMAMix', False):
+                        d_hat = min(self.d * (2 ** 0.1), d_hat)
+                    else:
+                        d_hat = min(self.d * (2 ** 0.25), d_hat)
                 if self.d == g_group['d0']:
                     self.d = max(self.d, d_hat)
                 d_max = max(d_max, d_hat)

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm/optim/Simplified_AdEMAMix.py

@@ -67,7 +67,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
             logging of Kourkoutas-β statistics (min, max, mean of `β₂` across layers)
             every logging steps. Useful for debugging and tuning. Set to 0 to disable
             logging (default: 0). 
-        layer_key_kb_fn (Optional[Callable]): A function that takes a parameter `p`
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
             and returns a unique, hashable key representing its "layer" or "bucket".
             If `None`, parameters are bucketed by their memory ID (tensor-wise).
             (default: None)
@@ -95,7 +95,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
         tiny_spike: float = 1e-9,
         k_warmup_steps: int = 0,
         k_logging: int = 0,
-        layer_key_kb_fn: Optional[Callable] = None,
+        layer_key_fn: Optional[Callable] = None,
         nnmf_factor: bool = False,
     ):
         if not (lr >= 0.0):
@@ -121,7 +121,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
         self.stochastic_rounding = stochastic_rounding
         self.factored = nnmf_factor
         self.kourkoutas_beta = kourkoutas_beta
-        self.layer_key_kb_fn = layer_key_kb_fn
+        self.layer_key_fn = layer_key_fn
         super().__init__(params, defaults)
 
         if self.kourkoutas_beta:

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm/util/Kourkoutas.py

@@ -32,21 +32,24 @@ class KourkoutasHelper:
         if self._layer_info_built:
             return
 
-        if hasattr(self.optimizer, 'layer_key_kb_fn') and self.optimizer.layer_key_kb_fn is not None:
+        if hasattr(self.optimizer, 'layer_key_fn') and self.optimizer.layer_key_fn is not None:
             # A custom key function was provided by the user. We will use it.
             pass
         else:
             # No key function was provided. Default to coarse, shape-based bucketing.
-            self.optimizer.layer_key_kb_fn = lambda p: \
+            self.optimizer.layer_key_fn = lambda p: \
                 (id(p),) if p.dim() == 2 and 1 <= p.shape[0] <= 10 and p.shape[1] in {768, 1280, 4096} \
                 else tuple(p.shape)
             # This ensures that we won't mix embeddings with tokens (1 to 10)
             # TODO find a better way to safeguard the embeddings
 
         for group in self.optimizer.param_groups:
+            if not group.get('kourkoutas_beta', False):
+                continue
+
             for p in group['params']:
                 # The mapping is static and should not depend on the presence of a gradient.
-                layer_key = self.optimizer.layer_key_kb_fn(p)
+                layer_key = self.optimizer.layer_key_fn(p)
                 if layer_key not in self.layer_info:
                     self.layer_info[layer_key] = {'params': [], 'group_ref': group}
                 self.layer_info[layer_key]['params'].append(p)
@@ -94,7 +97,7 @@ class KourkoutasHelper:
         for layer_key, info in self.layer_info.items():
             params, group = info['params'], info['group_ref']
 
-            if not group.get('kourkoutas_beta', False) and not group.get('_kourkoutas_beta', False):
+            if not group.get('kourkoutas_beta', False):
                 continue
 
             first_param_in_layer = info['params'][0]
@@ -174,7 +177,7 @@ class KourkoutasHelper:
         """
         Accumulates the squared L2 norm of a single gradient for the next step's calculation.
         """
-        layer_key = self.optimizer.layer_key_kb_fn(p)
+        layer_key = self.optimizer.layer_key_fn(p)
 
         if layer_key in self.layer_info:
             # Initialize the transient state for this layer if it's the first time in the step.
@@ -189,6 +192,6 @@ class KourkoutasHelper:
         """
         Gets the appropriate beta2 for the current parameter, handling warmup and dynamic value fetching.
         """
-        layer_key = self.optimizer.layer_key_kb_fn(p)
+        layer_key = self.optimizer.layer_key_fn(p)
         # The default is the max value, which is correct for unmapped params or edge cases
         return self.layer_state.get(layer_key, {}).get('dynamic_beta2', group['betas'][1])

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm/util/__init__.py

@@ -3,6 +3,7 @@ from .Effective_Shape import _get_effective_shape
 from .One_Bit_Boolean import _pack_bools, _unpack_bools
 from .OrthoGrad import _orthogonalize_gradient
 from .Newton_Schulz import _newton_schulz_iteration
+
 __all__ = [
     "_pack_bools", "_unpack_bools",
     "add_stochastic_",

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-1.2.dev8 → adv_optm-1.2.dev10}/adv_optm.egg-info/SOURCES.txt

@@ -19,7 +19,6 @@ 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.dev8 → adv_optm-1.2.dev10}/setup.py

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

adv_optm-1.2.dev8/adv_optm/util/MuonAdam_helper.py

@@ -1,32 +0,0 @@
-import torch
-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'