adv-optm 2.2.dev1__tar.gz → 2.2.dev3__tar.gz

{adv_optm-2.2.dev1 → adv_optm-2.2.dev3}/PKG-INFO

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

{adv_optm-2.2.dev1 → adv_optm-2.2.dev3}/adv_optm/__init__.py

@@ -22,4 +22,4 @@ __all__ = [
     "SignSGD_adv",
 ]
 
-__version__ = "2.2.dev1"
+__version__ = "2.2.dev3"

{adv_optm-2.2.dev1 → adv_optm-2.2.dev3}/adv_optm/optim/AdaMuon_adv.py

@@ -225,6 +225,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             "adam_k_warmup_steps": adam_k_warmup_steps, "adam_nnmf_factor": adam_nnmf_factor,
         }
         self.stochastic_rounding = stochastic_rounding
+        self._init_lr = lr
 
         super().__init__(params, defaults)
 
@@ -439,9 +440,13 @@ class AdaMuon_adv(torch.optim.Optimizer):
             scaled_eps, adaptive_eps, spectral_target, wd_scale = get_spectral_scaling(shape_for_scaling, group['n_layers'])
 
             weight_decay = group['weight_decay'] * wd_scale
+            decoupled_wd = True
+
             ns_eps = scaled_eps
+
         else:
             weight_decay = group['weight_decay']
+            decoupled_wd = False
             ns_eps = group['ns_eps']
             adaptive_eps = group['eps']
 
@@ -587,7 +592,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
             update = update.reshape(original_shape)
 
-        param_update.apply_parameter_update(self, p, group, update, lr, wd=weight_decay, random_int_tensor=random_int_tensor)
+        param_update.apply_parameter_update(self, p, group, update, lr, wd=weight_decay, random_int_tensor=random_int_tensor, decoupled=decoupled_wd)
 
     @torch.no_grad()
     def step(self, closure=None):

{adv_optm-2.2.dev1 → adv_optm-2.2.dev3}/adv_optm/optim/Muon_adv.py

@@ -204,6 +204,7 @@ class Muon_adv(torch.optim.Optimizer):
         }
         self.stochastic_rounding = stochastic_rounding
         self.compiled_optimizer = compiled_optimizer
+        self._init_lr = lr
 
         super().__init__(params, defaults)
 
@@ -404,9 +405,12 @@ class Muon_adv(torch.optim.Optimizer):
             scaled_eps, _, spectral_target, wd_scale = get_spectral_scaling(shape_for_scaling, group['n_layers'])
 
             weight_decay = group['weight_decay'] * wd_scale
+            decoupled_wd = True
+
             ns_eps = scaled_eps
         else:
             weight_decay = group['weight_decay']
+            decoupled_wd = False
             ns_eps = group['ns_eps']
 
         # MARS-M Approximated (Variance Reduction)
@@ -521,7 +525,7 @@ class Muon_adv(torch.optim.Optimizer):
 
                 update = update.reshape(original_shape)
 
-        param_update.apply_parameter_update(self, p, group, update, lr, wd=weight_decay, random_int_tensor=random_int_tensor)
+        param_update.apply_parameter_update(self, p, group, update, lr, wd=weight_decay, random_int_tensor=random_int_tensor, decoupled=decoupled_wd)
 
     @torch.no_grad()
     def step(self, closure=None):

{adv_optm-2.2.dev1 → adv_optm-2.2.dev3}/adv_optm/util/Muon_util.py

@@ -352,7 +352,12 @@ def spectral_norm_update(update: torch.Tensor, vector_state: torch.Tensor, targe
     # Normalize v_new to get next state
     v_norm = torch.linalg.vector_norm(v_new)
 
-    vector_state.copy_(v_new.div_(v_norm.clamp_min_(1e-12)).to(vector_state.dtype))
+    # if v_norm >= 0.5:
+    #    vector_state.copy_(v_new.div_(v_norm.clamp_min_(1e-12))).to(vector_state.dtype))
+    candidate_v = v_new / v_norm
+    next_state = torch.where(v_norm >= 0.5, candidate_v, vector_state)
+    vector_state.copy_(next_state.to(vector_state.dtype))
+    # Else: We keep the old vector_state (which is a random unit vector at init)
 
     # Estimate sigma = ||A @ v|| (since v is unit norm)
     # Re-compute A @ v_new with the updated vector for better estimate
@@ -379,7 +384,7 @@ def get_spectral_scaling(shape: torch.Size, n_layers: int):
         wd_scale: Weight decay scale
     """
     d_out, d_in = shape[0], shape[1]
-    
+
     # Handle Convolutional/Flattened tensors
     if len(shape) > 2:
         d_in = shape[1:].numel()

{adv_optm-2.2.dev1 → adv_optm-2.2.dev3}/adv_optm/util/param_update.py

@@ -14,6 +14,7 @@ def apply_parameter_update(
     lr: float | Tensor,
     wd: float | None = None,
     random_int_tensor: Tensor | None = None,
+    decoupled: bool = False,
 ) -> None:
     """
     Applies decoupled weight decay (standard or cautious) and the final
@@ -27,9 +28,14 @@ def apply_parameter_update(
         wd: Optional float value for weight decay, if another value other than group["weight_decay"] is needed.
         random_int_tensor: Optional pre-generated random tensor for stochastic
             rounding. Required for the `torch.compile` path.
+        decoupled: Whenever to use the true decoupled weight decay.
     """
     wd = group["weight_decay"] if wd is None else wd
     cautious = group.get('cautious_wd', False)
+    if decoupled:
+        scaled_wd = 1 / (1.0 +  wd * lr / self._init_lr)
+    else:
+        scaled_wd = wd * lr
 
     # Compute full update in float32 if using bfloat16 with stochastic rounding
     if p.dtype == torch.bfloat16 and self.stochastic_rounding:
@@ -41,11 +47,11 @@ def apply_parameter_update(
             if cautious:
                 # Cautious Weight Decay
                 mask = (update_fp32 * p_fp32 >= 0).float()
-                p_fp32.addcmul_(p_fp32, mask, value=-wd * lr)
+                p_fp32.addcmul_(p_fp32, mask, value=-scaled_wd)
                 del mask
             else:
                 # Standard decoupled weight decay
-                p_fp32.add_(p_fp32, alpha=-wd * lr)
+                p_fp32.add_(p_fp32, alpha=-scaled_wd)
 
         # Apply main update
         p_fp32.add_(-update_fp32)
@@ -65,11 +71,11 @@ def apply_parameter_update(
             if cautious:
                 # Cautious Weight Decay
                 mask = (update * p >= 0).to(p.dtype)
-                p.addcmul_(p, mask, value=-wd * lr)
+                p.addcmul_(p, mask, value=-scaled_wd)
                 del mask
             else:
                 # Standard decoupled weight decay
-                p.add_(p, alpha=-wd * lr)
+                p.add_(p, alpha=-scaled_wd)
 
         # Apply main update
         p.add_(-update)

{adv_optm-2.2.dev1 → adv_optm-2.2.dev3}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-2.2.dev1 → adv_optm-2.2.dev3}/setup.py

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