adv-optm 1.2.dev9__py3-none-any.whl → 1.2.dev11__py3-none-any.whl

adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "1.2.dev9"
+__version__ = "1.2.dev11"

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
@@ -12,7 +11,7 @@ from ..util.One_Bit_Boolean import _pack_bools, _unpack_bools
 
 class AdaMuon_adv(torch.optim.Optimizer):
     """
-    Implements the AdaMuon optimizer algorithm.
+    IImplements an advanced AdaMuon optimizer algorithm.
 
     AdaMuon combines the geometry-aware updates of Muon with the element-wise
     adaptivity of Adam. It is designed for 2D parameters (e.g., linear layers)
@@ -26,9 +25,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
     3.  An RMS-aligned rescaling strategy to match the update magnitude of Adam,
         allowing for reuse of learning rate schedules.
 
-    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
@@ -70,16 +66,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
             (default: 128)
         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__(
@@ -104,11 +90,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
         low_rank_ortho: bool = False,
         ortho_rank: int = 128,
         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}")
@@ -120,7 +101,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             print("Warning: nesterov is incompatible with Simplified_AdEMAMix, Disabling cautious.")
             nesterov = False
 
-        muon_defaults = {
+        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,
@@ -132,44 +113,8 @@ 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
 
-        # 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, muon_defaults)
-        
-        # Now that self is initialized, create the helper
-        self.helper = MuonAdamHelper(self, layer_key_fn)
+        super().__init__(params, defaults)
         
 
     @property
@@ -184,29 +129,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
     def supports_flat_params(self):
         return False
 
-    @property
-    def kourkoutas_helper(self):
-        """
-        Exposes the kourkoutas_helper from the auxiliary AdamW optimizer,
-        if it exists. This allows external access for logging K-b.
-        """
-        if self.aux_adam and hasattr(self.aux_adam, 'kourkoutas_helper'):
-            return self.aux_adam.kourkoutas_helper
-        return None
-
     @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
-                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
 

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,9 +105,8 @@ 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,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -137,12 +136,11 @@ 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
-        if not _is_delegate:
-            super().__init__(params, defaults)
-        else:
-            self.defaults = defaults
-        self.kourkoutas_helper = None
+        self.layer_key_fn = layer_key_fn
+        super().__init__(params, defaults)
+
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
 
     @property
     def supports_fused_back_pass(self):
@@ -160,8 +158,6 @@ 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:
@@ -244,6 +240,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 +250,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/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/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
@@ -25,10 +23,6 @@ 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.
-    
-    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
@@ -71,16 +65,6 @@ class Muon_adv(torch.optim.Optimizer):
         normuon_lr_scale (float): Scaling factor for the NorMuon learning rate.
             (default: 0.2)
         normuon_atan2 (bool): whether to use the atan2 for NorMuon. (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__(
@@ -108,11 +92,6 @@ class Muon_adv(torch.optim.Optimizer):
         normuon_eps: float = 1e-8,
         normuon_lr_scale: float = 0.2,
         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,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -128,7 +107,7 @@ class Muon_adv(torch.optim.Optimizer):
             print("Warning: nesterov is incompatible with Simplified_AdEMAMix, Disabling cautious.")
             nesterov = False
 
-        muon_defaults = {
+        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,
@@ -144,44 +123,7 @@ 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
-
-        # 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, muon_defaults)
-        
-        # Now that self is initialized, create the helper
-        self.helper = MuonAdamHelper(self, layer_key_fn)
+        super().__init__(params, defaults)
 
 
     @property
@@ -196,30 +138,8 @@ class Muon_adv(torch.optim.Optimizer):
     def supports_flat_params(self):
         return False
 
-    @property
-    def kourkoutas_helper(self):
-        """
-        Exposes the kourkoutas_helper from the auxiliary AdamW optimizer,
-        if it exists. This allows external access for logging K-b.
-        """
-        if self.aux_adam and hasattr(self.aux_adam, 'kourkoutas_helper'):
-            return self.aux_adam.kourkoutas_helper
-        return None
-
     @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
 

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/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/util/Kourkoutas.py

@@ -24,57 +24,35 @@ class KourkoutasHelper:
         # making it compatible with fused back pass mechanisms.
         self._build_layer_info_if_needed()
 
-        if self.optimizer.param_groups[0].get('k_logging', 0) > 0:
-            self.print_layer_info()
-
     def _build_layer_info_if_needed(self):
         """Builds a map of layers and the parameters they contain."""
         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)
-        
-        k_logging_interval = self.optimizer.param_groups[0].get('k_logging', 0)
-        if k_logging_interval > 0:
-            print(f"[Kourkoutas-β Debug] Layer info built. Found {len(self.layer_info)} unique layers/buckets.")
 
         self._layer_info_built = True
 
-    def print_layer_info(self):
-        """Prints the contents of self.layer_info for debugging."""
-        print("\n--- BEGIN self.layer_info DUMP ---")
-        if not self.layer_info:
-            print("Layer info is empty. Make sure the optimizer has parameters.")
-            return
-
-        for layer_key, info in self.layer_info.items():
-            param_count = len(info['params'])
-            first_param_details = ""
-            if param_count > 0:
-                p = info['params'][0]
-                first_param_details = f" (Example param shape: {list(p.shape)}, dtype: {p.dtype})"
-            
-            print(f"Key: {layer_key}, Params: {param_count}{first_param_details}")
-
-        print("--- END self.layer_info DUMP ---\n")
-
     def prepare_step(self, current_step: int):
         """
         Calculates dynamic beta2 for all layers using the completed scalar accumulators
@@ -82,9 +60,8 @@ class KourkoutasHelper:
         """
         
         beta2_log = []
-        first_layer_key = next(iter(self.layer_info), None)
         # These are just for the sample log, initialize them
-        sun, pooled_grad_norm, prev_r_ema_val, r_ema_tensor = (torch.tensor(0.0),)*4
+        sun, pooled_grad_norm, r_ema_tensor = (torch.tensor(0.0),)*3
 
         # The optimizer that owns this helper holds the master defaults for K-b.
         # This is crucial in hybrid optimizers where some param_groups might not
@@ -94,7 +71,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]
@@ -121,7 +98,6 @@ class KourkoutasHelper:
             accumulator = self.layer_state[layer_key]['sum_sq_accumulator']
             
             pooled_grad_norm = torch.sqrt(accumulator)
-            prev_r_ema_val = r_ema_tensor.item() # for logging
             
             # Update the persistent EMA tensor in-place.
             r_ema_tensor.mul_(ema_alpha).add_(pooled_grad_norm, alpha=1.0 - ema_alpha)
@@ -147,21 +123,9 @@ class KourkoutasHelper:
         if beta2_log:
             beta2_tensor = torch.tensor(beta2_log, device='cpu')
             self.last_beta2_stats = {
-                'min': beta2_tensor.min().item(),
-                'max': beta2_tensor.max().item(),
                 'mean': beta2_tensor.mean().item(),
             }
 
-        # Handle periodic console logging
-        k_logging_interval = self.optimizer.param_groups[0].get('k_logging', 0)
-        is_logging_step = k_logging_interval > 0 and (current_step + 1) % k_logging_interval == 0
-        if is_logging_step and self.last_beta2_stats:
-            if first_layer_key:
-                print(f"\n[Kourkoutas-β Debug] Step {current_step + 1} - Sample Layer '{first_layer_key}':")
-                print(f"  - Grad Norm: {pooled_grad_norm.item():.4e}, Prev EMA: {prev_r_ema_val:.4e}, New EMA: {r_ema_tensor.item():.4e}")
-                print(f"  - Sunspike: {sun.item():.4f}, Dynamic Beta2: {self.layer_state[first_layer_key]['dynamic_beta2']:.4f}")
-            print(f"[Kourkoutas-β Debug] Step {current_step + 1} Overall Beta2 Stats: Min={self.last_beta2_stats['min']:.4f}, Max={self.last_beta2_stats['max']:.4f}, Mean={self.last_beta2_stats['mean']:.4f}")
-
     def maybe_prepare_step(self, current_step: int):
         """
         A universal guard that calls prepare_step() exactly once per training step.
@@ -174,7 +138,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 +153,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/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.dev9.dist-info → adv_optm-1.2.dev11.dist-info}/METADATA

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

adv_optm-1.2.dev11.dist-info/RECORD

@@ -0,0 +1,23 @@
+adv_optm/__init__.py,sha256=vjm5Sc3hgTSy9qP73qVBVGp9zE2J6blsQmj6KxyI3GE,380
+adv_optm/optim/AdaMuon_adv.py,sha256=828WtdsaKXJqlZqFXE2yrsxY3Erxn-6N7CxV9jBXiaI,17880
+adv_optm/optim/AdamW_adv.py,sha256=KL9SCJWZ_ckAQEApB6ofbndVYjancN-v7Us7hJLFf54,17475
+adv_optm/optim/Adopt_adv.py,sha256=S8XI2YA7683jsW8p7igc2YcU30lsN0H18qL02Kpvj8E,21244
+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=xGW9PafaIyi1noGhIgCWPwndI5bGX6kbxN-N-FQnr1U,19381
+adv_optm/optim/Prodigy_adv.py,sha256=lEjbtuQbomsCX39DnTPeI8Z5YG0f2aZPXN_E7-nGgWw,26060
+adv_optm/optim/Simplified_AdEMAMix.py,sha256=nEIA3yM11nBooKzHudB5l3x4UdFRBYRwiKVUkGmO0K8,12971
+adv_optm/optim/__init__.py,sha256=hpUWE6CKtt_rvMdgQVb3PtjhfZAvAxTq6hp8H8rIpBo,489
+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=_fq2glPqKmzgWpLedfwq5EqIJAxICUK2fmUP-cdcgq0,7467
+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=CXzS703GB4gil85khZi7sgKOnbzXGBOltshIOSPqj18,435
+adv_optm-1.2.dev11.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-1.2.dev11.dist-info/METADATA,sha256=F6o4bbgIEjB9JS_9gediI_0-_rUkfsubKVtg5b4nrHE,14023
+adv_optm-1.2.dev11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+adv_optm-1.2.dev11.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-1.2.dev11.dist-info/RECORD,,

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'

adv_optm-1.2.dev9.dist-info/RECORD

@@ -1,24 +0,0 @@
-adv_optm/__init__.py,sha256=TzvKgGTLkK0_XANeZzhURcSO9xmtUi-H9_C7tV3rXn4,379
-adv_optm/optim/AdaMuon_adv.py,sha256=yr1oJV339Zv7D8n148O1FJJAgdOsH8NZDZTKlcDOyu0,21181
-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=HaF06fPKcKpVZY29_vqjWHAfivjvGntBuRyDDKj3Ozw,22784
-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
-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=lObJGXmz3MqGSuu3DKqotSpZ0fuQFPE80R3zO_j3Z_Q,9707
-adv_optm/util/MuonAdam_helper.py,sha256=7rnNMujZVDaqo1g22QscMyPlZvIHQQSLHMED9_I8QWU,1250
-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.dev9.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-1.2.dev9.dist-info/METADATA,sha256=GmAYWjZdfgvg9QbzyiV2PUNmzQFgJz8AjaY5F0x7Nv8,14022
-adv_optm-1.2.dev9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-1.2.dev9.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-1.2.dev9.dist-info/RECORD,,