adv-optm 1.2.dev17__py3-none-any.whl → 1.2.dev19__py3-none-any.whl

adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "1.2.dev17"
+__version__ = "1.2.dev19"

adv_optm/optim/AdaMuon_adv.py

@@ -46,6 +46,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             (default: (3.4445, -4.7750, 2.0315)).
         stochastic_rounding (bool): whether to use stochastic rounding for
             BF16 parameter updates (default: True).
+        orthogonal_gradient (bool): whether to use OrthoGrad.  (default: False)
         nesterov (bool): enables Nesterov momentum (default: False).
         use_atan2 (bool): whether to use the atan2 update rule. (default: False)
         Simplified_AdEMAMix (bool): whether to use the Simplified AdEMAMix update rule.
@@ -95,6 +96,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         ns_eps: float = 1e-7,
         ns_coeffs: tuple[float, float, float] = (3.4445, -4.7750, 2.0315),
         stochastic_rounding: bool = False,
+        orthogonal_gradient: bool = False,
         use_atan2: bool = False,
         nesterov: bool = False,
         Simplified_AdEMAMix: bool = False,
@@ -147,7 +149,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             "vector_reshape": vector_reshape,
             "nesterov":nesterov, "use_atan2":use_atan2,
             "Simplified_AdEMAMix": Simplified_AdEMAMix, "alpha_grad": alpha_grad,
-            "normuon_variant": normuon_variant,
+            "normuon_variant": normuon_variant, "orthogonal_gradient": orthogonal_gradient,
             # Low-rank Ortho
             "low_rank_ortho": low_rank_ortho, "ortho_rank": ortho_rank,
             "compiled_optimizer":compiled_optimizer,
@@ -282,6 +284,10 @@ class AdaMuon_adv(torch.optim.Optimizer):
         nesterov = group['nesterov']
         Simplified_AdEMAMix = group['Simplified_AdEMAMix']
         alpha_grad = group['alpha_grad']
+        if grad.dtype != torch.float32 and state.get('factored', False):
+            grad = grad.float()
+        if group.get("orthogonal_gradient"):
+            grad = _orthogonalize_gradient(p, grad)
 
         if state['factored']: # Factored AdaMuon
 
@@ -345,11 +351,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 mean_squared_update = torch.mean(update.square(), dim=1)
                 v_t.mul_(beta2).add_(mean_squared_update, alpha=1 - beta2)
                 # Normalize update
-                if group['use_atan2']:
-                    a = 1.2732395
-                    update.atan2_(v_t.sqrt().unsqueeze(1)).mul_(a)
-                else:
-                    update.div_(v_t.sqrt().unsqueeze(1).add_(group['eps']))
+                update.div_(v_t.sqrt().unsqueeze(1).add_(group['eps']))
                 # Scale learning rate
                 update_norm = torch.linalg.vector_norm(update)
                 scaled_lr = group['rms_target'] * lr * (p.numel()**0.5) / update_norm.add_(group['eps'])
@@ -454,11 +456,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
                     mean_squared_update = torch.mean(update.square(), dim=1)
                     v_t.mul_(beta2).add_(mean_squared_update, alpha=1 - beta2)
                     # Normalize update
-                    if group['use_atan2']:
-                        a = 1.2732395
-                        update.atan2_(v_t.sqrt().unsqueeze(1)).mul_(a)
-                    else:
-                        update.div_(v_t.sqrt().unsqueeze(1).add_(group['eps']))
+                    update.div_(v_t.sqrt().unsqueeze(1).add_(group['eps']))
                     # Scale learning rate
                     update_norm = torch.linalg.vector_norm(update)
                     scaled_lr = group['rms_target'] * lr * (p.numel()**0.5) / update_norm.add_(group['eps'])

adv_optm/optim/AdamW_adv.py

@@ -252,7 +252,7 @@ class AdamW_adv(torch.optim.Optimizer):
                 # Update momentum in full-size
                 mt.mul_(beta1).add_(grad_reshaped, alpha=1.0 - beta1)
                 if self.grams_moment:
-                    mt.copy_(grad_reshaped.sign() * mt.abs())
+                    mt = (grad_reshaped.sign().mul_(mt.abs()))
                 elif self.cautious_mask:
                     mask = (mt * grad_reshaped > 0).to(grad_reshaped.dtype)
                     mask.div_(mask.mean().clamp_(min=1e-3))
@@ -310,7 +310,7 @@ class AdamW_adv(torch.optim.Optimizer):
                 exp_avg = state['exp_avg']
                 exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                 if self.grams_moment:
-                    exp_avg = grad.sign() * exp_avg.abs()
+                    exp_avg = grad.sign().mul_(exp_avg.abs())
                 elif self.cautious_mask:
                     mask = (exp_avg * grad > 0).to(grad.dtype)
                     mask.div_(mask.mean().clamp_(min=1e-3))

adv_optm/optim/Adopt_adv.py

@@ -307,7 +307,7 @@ class Adopt_adv(torch.optim.Optimizer):
                 else:
                     mt.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
                 if self.grams_moment:
-                    mt = grad_reshaped.sign() * mt.abs()
+                    mt = grad_reshaped.sign().mul_(mt.abs())
                 elif self.cautious_mask:
                     mask = (mt * grad_reshaped > 0).to(grad_reshaped.dtype)
                     mask.div_(mask.mean().clamp_(min=1e-3))
@@ -376,7 +376,7 @@ class Adopt_adv(torch.optim.Optimizer):
                     m.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
 
             if self.grams_moment:
-                m = grad.sign() * m.abs()
+                m = grad.sign().mul_(m.abs())
             elif self.cautious_mask:
                 mask = (m * grad > 0).to(grad.dtype)
                 mask.div_(mask.mean().clamp_(min=1e-3))

adv_optm/optim/Muon_adv.py

@@ -41,6 +41,7 @@ class Muon_adv(torch.optim.Optimizer):
             stability. (default: 100.0)
         stochastic_rounding (bool): whether to use stochastic rounding for
             BF16 parameter updates (default: True).
+        orthogonal_gradient (bool): whether to use OrthoGrad.  (default: False)
         vector_reshape_muon (bool): whether to reshape 1D vectors into 2D
             matrices for muon NewtonSchulz (default: False).
         vector_reshape (bool): whether to reshape 1D vectors into 2D
@@ -59,7 +60,6 @@ class Muon_adv(torch.optim.Optimizer):
         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)
         accelerated_ns (bool): If True, enables Chebyshev-accelerated Newton-Schulz, which
             dynamically calculates optimal 3rd-order polynomial coefficients. (default: False)
         cns_a_bound (float): Initial lower bound for singular values for CANS. (default: 1e-4)
@@ -92,6 +92,7 @@ class Muon_adv(torch.optim.Optimizer):
         Simplified_AdEMAMix: bool = False,
         alpha_grad: float = 100.0,
         stochastic_rounding: bool = True,
+        orthogonal_gradient: bool = False,
         vector_reshape_muon: bool = False,
         vector_reshape: bool = False,
         nnmf_factor: bool = False,
@@ -103,7 +104,6 @@ class Muon_adv(torch.optim.Optimizer):
         beta2_normuon: float = 0.95,
         normuon_eps: float = 1e-8,
         normuon_lr_scale: float = 0.2,
-        normuon_atan2: bool = False,
         # CANS
         accelerated_ns: bool = False,
         cns_a_bound: float = 1e-4,
@@ -149,13 +149,13 @@ class Muon_adv(torch.optim.Optimizer):
             "vector_reshape": vector_reshape,
             "vector_reshape_muon": vector_reshape_muon,
             "Simplified_AdEMAMix": Simplified_AdEMAMix, "alpha_grad": alpha_grad,
+            "orthogonal_gradient": orthogonal_gradient,
             'compiled_optimizer': compiled_optimizer,
             # 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,
             # CANS
             "accelerated_ns": accelerated_ns, "cns_a_bound": cns_a_bound,
             # AdamW_adv defaults
@@ -293,6 +293,10 @@ class Muon_adv(torch.optim.Optimizer):
         nesterov = group['nesterov']
         Simplified_AdEMAMix = group['Simplified_AdEMAMix']
         alpha_grad = group['alpha_grad']
+        if grad.dtype != torch.float32 and state.get('factored', False):
+            grad = grad.float()
+        if group.get("orthogonal_gradient"):
+            grad = _orthogonalize_gradient(p, grad)
 
         if state['factored']: # Factored Muon
 
@@ -359,11 +363,7 @@ class Muon_adv(torch.optim.Optimizer):
                 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']))
+                update.div_(v_t.sqrt().unsqueeze(1).add_(group['normuon_eps']))
                 # Scale learning rate
                 update_norm = torch.linalg.vector_norm(update)
 
@@ -464,11 +464,7 @@ class Muon_adv(torch.optim.Optimizer):
                     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']))
+                    update.div_(v_t.sqrt().unsqueeze(1).add_(group['normuon_eps']))
                     # Scale learning rate
                     update_norm = torch.linalg.vector_norm(update)
                     scaled_lr = group['normuon_lr_scale'] * lr * (p.numel()**0.5) / update_norm.add_(group['normuon_eps'])

adv_optm/optim/Prodigy_adv.py

@@ -343,7 +343,7 @@ class Prodigy_adv(torch.optim.Optimizer):
                 else:
                     mt.mul_(self.beta1).add_(grad_reshaped, alpha=self.d * (1.0 - self.beta1))
                 if self.grams_moment:
-                    mt.copy_(grad_reshaped.sign() * mt.abs())
+                    mt = (grad_reshaped.sign().mul_(mt.abs()))
                 elif self.cautious_mask:
                     mask = (mt * grad_reshaped > 0).to(grad_reshaped.dtype)
                     mask.div_(mask.mean().clamp_(min=1e-3))

{adv_optm-1.2.dev17.dist-info → adv_optm-1.2.dev19.dist-info}/METADATA

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

{adv_optm-1.2.dev17.dist-info → adv_optm-1.2.dev19.dist-info}/RECORD

@@ -1,11 +1,11 @@
-adv_optm/__init__.py,sha256=5Dww3w78iQNwyVH82E_dmD-s6luvQjiqYS0BxKQHYCE,380
-adv_optm/optim/AdaMuon_adv.py,sha256=zjZHFS7ng5KwemQzePjFiGtNZlcgbzmmnqF6A80h_Tg,34652
-adv_optm/optim/AdamW_adv.py,sha256=KL9SCJWZ_ckAQEApB6ofbndVYjancN-v7Us7hJLFf54,17475
-adv_optm/optim/Adopt_adv.py,sha256=S8XI2YA7683jsW8p7igc2YcU30lsN0H18qL02Kpvj8E,21244
+adv_optm/__init__.py,sha256=1AKxG--scx5Bl9G08tQcnfzAMaQVSgmW99uy3v2QWMw,380
+adv_optm/optim/AdaMuon_adv.py,sha256=7Had92OcsCiN1E9UJRyrpPV7VzHqmIvS-qM6OEcc24I,34671
+adv_optm/optim/AdamW_adv.py,sha256=jgMuRAfsnUh_2wUEZgYpJX5uwoT_kQjtMs2Xn2vJ3x0,17480
+adv_optm/optim/Adopt_adv.py,sha256=kbAeBG4bXWBvgj_qrE9W67J6c0swpEi4Erj2rfYrMXE,21252
 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=d91wvmKKt_3IPqsqK1ZZ5cY71kuXyzy04IU3krn2NQ8,33316
-adv_optm/optim/Prodigy_adv.py,sha256=lEjbtuQbomsCX39DnTPeI8Z5YG0f2aZPXN_E7-nGgWw,26060
+adv_optm/optim/Muon_adv.py,sha256=tZY8K3pNBCGk1V09GbK05lJooFw92NfkF7_T548up3Q,33171
+adv_optm/optim/Prodigy_adv.py,sha256=k7f2J_RQpnrUXjwER_XOokISlQWpTSwGG-OL-bjMfBk,26061
 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
@@ -16,8 +16,8 @@ adv_optm/util/Newton_Schulz.py,sha256=bBboYw_jm5_FMf0Citl79uqNedkHOTjQnUI7rZgLBm
 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.dev17.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-1.2.dev17.dist-info/METADATA,sha256=xE_ECrY_ALerNQRFBtKml1w_n8wSp8zOH0tIz-BLiqY,14023
-adv_optm-1.2.dev17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-1.2.dev17.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-1.2.dev17.dist-info/RECORD,,
+adv_optm-1.2.dev19.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-1.2.dev19.dist-info/METADATA,sha256=pQm5WuMKvf5Xse10viziVK9ry1UufcYRDwOd55jad8Y,14023
+adv_optm-1.2.dev19.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+adv_optm-1.2.dev19.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-1.2.dev19.dist-info/RECORD,,