adv-optm 1.2.dev5__tar.gz → 1.2.dev7__tar.gz

{adv_optm-1.2.dev5 → adv_optm-1.2.dev7}/PKG-INFO

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

{adv_optm-1.2.dev5 → adv_optm-1.2.dev7}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "1.2.dev5"
+__version__ = "1.2.dev7"

{adv_optm-1.2.dev5 → adv_optm-1.2.dev7}/adv_optm/optim/Muon_adv.py

@@ -18,6 +18,10 @@ class Muon_adv(torch.optim.Optimizer):
     the hidden layers of neural networks. It applies SGD with momentum and then
     orthogonalizes the resulting update matrix using a Newton-Schulz iteration.
 
+    NorMuon (Neuron-wise Normalized Muon) extends this by adding neuron-level
+    adaptive learning rates, combining the benefits of orthogonalization with
+    second-order momentum statistics.
+
     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.
@@ -54,6 +58,19 @@ 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)
+        low_rank_ortho (bool): If True, enables low-rank orthogonalization, which
+            projects the update to a lower rank before orthogonalization.
+            (default: False)
+        ortho_rank (int): The rank for low-rank orthogonalization.
+            (default: 128)
+        normuon_variant (bool): If True, enables the NorMuon update rule, which adds
+            neuron-wise normalization. (default: False)
+        beta2_normuon (float): The exponential decay rate for the second moment estimates
+            used in NorMuon. (default: 0.95)
+        normuon_eps (float): Epsilon for NorMuon normalization stability. (default: 1e-8)
+        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)
@@ -82,6 +99,15 @@ class Muon_adv(torch.optim.Optimizer):
         vector_reshape_muon: bool = False,
         vector_reshape: bool = False,
         nnmf_factor: bool = False,
+        # Low-rank Muon
+        low_rank_ortho: bool = False,
+        ortho_rank: int = 128,
+        # NorMuon additions
+        normuon_variant: bool = False,
+        beta2_normuon: float = 0.95,
+        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,
@@ -92,6 +118,8 @@ class Muon_adv(torch.optim.Optimizer):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
         if not (0.0 <= beta1 < 1.0):
             raise ValueError(f"beta1 should be in [0.0, 1.0). Got {beta1}")
+        if normuon_variant and not (0.0 <= beta2_normuon < 1.0):
+            raise ValueError(f"beta2_normuon should be in [0.0, 1.0) for NorMuon. Got {beta2_normuon}")
         if not (weight_decay >= 0.0):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
         if not (ns_steps > 0):
@@ -106,10 +134,16 @@ class Muon_adv(torch.optim.Optimizer):
             "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor,
             "vector_reshape": vector_reshape,
             "vector_reshape_muon": vector_reshape_muon,
-            "Simplified_AdEMAMix": Simplified_AdEMAMix, "alpha_grad": alpha_grad, 
+            "Simplified_AdEMAMix": Simplified_AdEMAMix, "alpha_grad": alpha_grad,
+            # Low-rank Ortho
+            "low_rank_ortho": low_rank_ortho, "ortho_rank": ortho_rank,
+            # NorMuon
+            "normuon_variant": normuon_variant, "beta2_normuon": beta2_normuon,
+            "normuon_eps": normuon_eps, "normuon_lr_scale": normuon_lr_scale,
+            "normuon_atan2": normuon_atan2,
         }
         self.stochastic_rounding = stochastic_rounding
-        
+
         self.MuonWithAuxAdam = MuonWithAuxAdam
         self.helper = None
         self.aux_adam = None
@@ -223,6 +257,12 @@ class Muon_adv(torch.optim.Optimizer):
                 elif len(p.shape) == 1:
                     state['momentum_buffer'] = torch.zeros_like(p)
 
+            # NorMuon state initialization
+            if group['normuon_variant']:
+                if 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)
+
         beta1 = group['beta1']
         nesterov = group['nesterov']
         Simplified_AdEMAMix = group['Simplified_AdEMAMix']
@@ -251,14 +291,60 @@ class Muon_adv(torch.optim.Optimizer):
                 update = mt_buf.clone()
             del grad_reshaped
 
-            update = _newton_schulz_iteration(
-                update,
-                steps=group['ns_steps'],
-                eps=group['ns_eps'],
-                coeffs=group['ns_coeffs'],
-            )
+            # Orthogonalization step
+            if group['low_rank_muon']:
+                # Low-Rank Orthogonalization on the reconstructed matrix
+                M = update
+                r = min(group['low_rank_rank'], M.shape[0], M.shape[1])
+                if r > 0:
+                    G_sketch = torch.randn(M.shape[1], r, device=M.device, dtype=M.dtype)
+                    MG = M @ G_sketch
+                    if MG.dtype != torch.float32:
+                        MG_dtype = M.dtype
+                        Q, _ = torch.linalg.qr(MG.float())
+                        Q = Q.to(MG_dtype)
+                    else:
+                        Q, _ = torch.linalg.qr(MG)
+                    projected_M = Q.T @ M
+                    ortho_projected_M = _newton_schulz_iteration(
+                        projected_M, steps=group['ns_steps'], eps=group['ns_eps'], coeffs=group['ns_coeffs']
+                    )
+                    update = Q @ ortho_projected_M
+                else: # Fallback for invalid rank
+                    update = _newton_schulz_iteration(
+                        update, steps=group['ns_steps'], eps=group['ns_eps'], coeffs=group['ns_coeffs']
+                    )
+            else:
+                # Original full Newton-Schulz
+                update = _newton_schulz_iteration(
+                    update,
+                    steps=group['ns_steps'],
+                    eps=group['ns_eps'],
+                    coeffs=group['ns_coeffs'],
+                )
 
-            update = update.view(p.shape).mul_(group['lr'])
+
+            if group['normuon_variant'] and 'normuon_v' in state:
+                v_t = state['normuon_v']
+                beta2_normuon = group['beta2_normuon']
+                # Update 2nd moment estimate
+                mean_squared_update = torch.mean(update.square(), dim=1)
+                v_t.mul_(beta2_normuon).add_(mean_squared_update, alpha=1 - beta2_normuon)
+                # Normalize update
+                if group['normuon_atan2']:
+                    a = 1.2732395
+                    update.atan2_(v_t.sqrt().unsqueeze(1)).mul_(a)
+                else:
+                    update.div_(v_t.sqrt().unsqueeze(1).add_(group['normuon_eps']))
+                # Scale learning rate
+                update_norm = torch.linalg.vector_norm(update)
+                if update_norm > 1e-12:
+                    scaled_lr = group['normuon_lr_scale'] * group['lr'] * (p.numel()**0.5) / update_norm
+                else:
+                    scaled_lr = 0.0
+                update = update.view(p.shape).mul_(scaled_lr)
+            else: # Original Muon learning rate application
+                update = update.view(p.shape).mul_(group['lr'])
 
             state['sign'] = _pack_bools(mt_buf > 0)
             _nnmf(mt_buf.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
@@ -298,19 +384,80 @@ class Muon_adv(torch.optim.Optimizer):
                 if len(p.shape) > 2:
                     update = update.view(p.shape[0], -1)
 
-                # NewtonSchulz
-                update = _newton_schulz_iteration(
-                    update,
-                    steps=group['ns_steps'],
-                    eps=group['ns_eps'],
-                    coeffs=group['ns_coeffs'],
-                )
+                # Orthogonalization step
+                if group['low_rank_ortho']:
+                    # Low-Rank Orthogonalization based on Gaussian Sketching
+                    M = update
+                    r = min(group['ortho_rank'], M.shape[0], M.shape[1])
+
+                    if r > 0:
+                        # 1. Sketch the matrix
+                        G_sketch = torch.randn(M.shape[1], r, device=M.device, dtype=M.dtype)
+                        MG = M @ G_sketch
+                        
+                        # 2. QR decomposition to get orthogonal basis Q
+                        if MG.dtype != torch.float32:
+                            MG_dtype = M.dtype
+                            Q, _ = torch.linalg.qr(MG.float())
+                            Q = Q.to(MG_dtype)
+                        else:
+                            Q, _ = torch.linalg.qr(MG)
+                        
+                        # 3. Project M onto the basis
+                        projected_M = Q.T @ M
+                        
+                        # 4. Orthogonalize the smaller projected matrix
+                        ortho_projected_M = _newton_schulz_iteration(
+                            projected_M,
+                            steps=group['ns_steps'],
+                            eps=group['ns_eps'],
+                            coeffs=group['ns_coeffs'],
+                        )
+                        
+                        # 5. Project back to the original space
+                        update = Q @ ortho_projected_M
+                    else: # Fallback for invalid rank
+                        update = _newton_schulz_iteration(
+                            update,
+                            steps=group['ns_steps'],
+                            eps=group['ns_eps'],
+                            coeffs=group['ns_coeffs'],
+                        )
+                else:
+                    # Original NewtonSchulz
+                    update = _newton_schulz_iteration(
+                        update,
+                        steps=group['ns_steps'],
+                        eps=group['ns_eps'],
+                        coeffs=group['ns_coeffs'],
+                    )
+
+                # NorMuon Logic
+                if group['normuon_variant'] and 'normuon_v' in state:
+                    v_t = state['normuon_v']
+                    beta2_normuon = group['beta2_normuon']
+                    # Update 2nd moment estimate
+                    mean_squared_update = torch.mean(update.square(), dim=1)
+                    v_t.mul_(beta2_normuon).add_(mean_squared_update, alpha=1 - beta2_normuon)
+                    # Normalize update
+                    if group['normuon_atan2']:
+                        a = 1.2732395
+                        update.atan2_(v_t.sqrt().unsqueeze(1)).mul_(a)
+                    else:
+                        update.div_(v_t.sqrt().unsqueeze(1).add_(group['normuon_eps']))
+                    # Scale learning rate
+                    update_norm = torch.linalg.vector_norm(update)
+                    if update_norm > 1e-12:
+                        scaled_lr = group['normuon_lr_scale'] * group['lr'] * (p.numel()**0.5) / update_norm
+                    else:
+                        scaled_lr = 0.0
+                    update.mul_(scaled_lr)
+                else: # Original Muon learning rate application
+                    update.mul_(group['lr'])
 
                 # Reshape back to original if we flattened or reshaped
                 if len(p.shape) > 2 or state['reshaped_1d_muon']:
                     update = update.view(p.shape)
-
-                update.mul_(group['lr'])
             
             else: # Fallback to standard SGD with momentum for 1D params (biases, etc.) when not reshaped
                 # Momentum update

{adv_optm-1.2.dev5 → adv_optm-1.2.dev7}/adv_optm/util/Kourkoutas.py

@@ -86,9 +86,17 @@ class KourkoutasHelper:
         # 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
 
+        # 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
+        # have all K-b keys populated, preventing KeyErrors.
+        master_defaults = self.optimizer.defaults
+
         for layer_key, info in self.layer_info.items():
             params, group = info['params'], info['group_ref']
 
+            if not group.get('kourkoutas_beta', False):
+                continue
+
             first_param_in_layer = info['params'][0]
             param_state = self.optimizer.state[first_param_in_layer]
 
@@ -100,6 +108,15 @@ class KourkoutasHelper:
             if 'kourkoutas_r_ema' not in param_state:
                 param_state['kourkoutas_r_ema'] = torch.tensor(0.0, device=first_param_in_layer.device, dtype=torch.float32)
 
+            # Use group-specific K-b settings, falling back to the optimizer's master defaults.
+            # This makes the helper robust against param groups that enable kourkoutas_beta
+            # but are missing the other required hyperparameters.
+            ema_alpha = group.get('ema_alpha', master_defaults['ema_alpha'])
+            beta2_max = group.get('betas', master_defaults['betas'])[1]
+            beta2_min = group.get('beta2_min', master_defaults['beta2_min'])
+            tiny_spike = group.get('tiny_spike', master_defaults['tiny_spike'])
+            k_warmup_steps = group.get('k_warmup_steps', master_defaults['k_warmup_steps'])
+
             r_ema_tensor = param_state['kourkoutas_r_ema']
             accumulator = self.layer_state[layer_key]['sum_sq_accumulator']
             
@@ -107,17 +124,16 @@ class KourkoutasHelper:
             prev_r_ema_val = r_ema_tensor.item() # for logging
             
             # Update the persistent EMA tensor in-place.
-            r_ema_tensor.mul_(group['ema_alpha']).add_(pooled_grad_norm, alpha=1.0 - group['ema_alpha'])
+            r_ema_tensor.mul_(ema_alpha).add_(pooled_grad_norm, alpha=1.0 - ema_alpha)
             
-            beta2_max = group['betas'][1]
             sun = torch.tensor(0.0, device=r_ema_tensor.device) # Default sun to 0 for warmup
             
-            if current_step < group['k_warmup_steps']:
+            if current_step < k_warmup_steps:
                 beta2 = beta2_max
             else:
-                raw = pooled_grad_norm / (r_ema_tensor + group['tiny_spike'])
+                raw = pooled_grad_norm / (r_ema_tensor + tiny_spike)
                 sun = raw / (1.0 + raw)
-                beta2 = beta2_max - (beta2_max - group['beta2_min']) * sun
+                beta2 = beta2_max - (beta2_max - beta2_min) * sun
 
             # Store the final calculated beta2 in the helper's transient state for this step.
             self.layer_state[layer_key]['dynamic_beta2'] = beta2.item() if isinstance(beta2, torch.Tensor) else beta2

{adv_optm-1.2.dev5 → adv_optm-1.2.dev7}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-1.2.dev5 → adv_optm-1.2.dev7}/setup.py

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