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

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

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

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

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

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

@@ -11,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)
@@ -25,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
@@ -69,12 +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)
-        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__(
@@ -99,10 +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,
-        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}")
@@ -114,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,
@@ -127,34 +114,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         }
         self.stochastic_rounding = stochastic_rounding
 
-        self.MuonWithAuxAdam = MuonWithAuxAdam
-        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 = []
-        for group in params:
-            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)
+        super().__init__(params, defaults)
         
 
     @property
@@ -169,30 +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 = 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
-
         if p.grad is None:
             return
 

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

@@ -107,7 +107,6 @@ 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}")
@@ -138,11 +137,10 @@ class AdamW_adv(torch.optim.Optimizer):
         self.factored = nnmf_factor
         self.kourkoutas_beta = kourkoutas_beta
         self.layer_key_fn = layer_key_fn
-        if not _is_delegate:
-            super().__init__(params, defaults)
-        else:
-            self.defaults = defaults
-        self.kourkoutas_helper = None
+        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:

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

@@ -23,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
@@ -69,12 +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)
-        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__(
@@ -102,10 +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,
-        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}")
@@ -121,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,
@@ -137,34 +123,7 @@ class Muon_adv(torch.optim.Optimizer):
         }
         self.stochastic_rounding = stochastic_rounding
 
-        self.MuonWithAuxAdam = MuonWithAuxAdam
-        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 = []
-        for group in params:
-            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)
+        super().__init__(params, defaults)
 
 
     @property
@@ -179,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 = 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
-
         if p.grad is None:
             return
 
@@ -242,7 +179,9 @@ class Muon_adv(torch.optim.Optimizer):
 
             # NorMuon state initialization
             if group['normuon_variant']:
-                if len(p.shape) >= 2 or state['reshaped_1d_muon']:
+                if state['factored']:
+                    state['normuon_v'] = torch.zeros(d1, device=p.device, dtype=torch.float32)
+                elif len(p.shape) >= 2 or state['reshaped_1d_muon']:
                     num_rows = p.shape[0] if len(p.shape) >= 2 else state['effective_shape'][0]
                     state['normuon_v'] = torch.zeros(num_rows, device=p.device, dtype=torch.float32)
 

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

@@ -24,9 +24,6 @@ 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:
@@ -53,31 +50,9 @@ class KourkoutasHelper:
                 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
@@ -85,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
@@ -124,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)
@@ -150,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.

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

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

{adv_optm-1.2.dev10 → adv_optm-1.2.dev12}/setup.py

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