adv-optm 1.0.5__py3-none-any.whl → 1.0.6__py3-none-any.whl

adv_optm/__init__.py

@@ -16,4 +16,4 @@ __all__ = [
     "Lion_Prodigy_adv",
 ]
 
-__version__ = "1.0.5"
+__version__ = "1.0.6"

adv_optm/optim/Adopt_adv.py

@@ -156,6 +156,8 @@ class Adopt_adv(torch.optim.Optimizer):
             grad = _orthogonalize_gradient(p, grad)
         state = self.state[p]
 
+        beta1, beta2 = group['betas']
+
         # State Initialization
         if len(state) == 0:
             state['step'] = 0
@@ -174,11 +176,12 @@ class Adopt_adv(torch.optim.Optimizer):
                 d1, d2 = state['effective_shape']
 
                 # m_0 = 0
-                state['mu_m_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype) 
-                state['mv_m_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
-                if not self.grams_moment:
-                    packed_d2 = (d2 + 7) // 8
-                    state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
+                if beta1 > 0:
+                    state['mu_m_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype) 
+                    state['mv_m_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
+                    if not self.grams_moment:
+                        packed_d2 = (d2 + 7) // 8
+                        state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
                 if self.use_AdEMAMix:
                     state['mu_m_slow_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype) 
                     state['mv_m_slow_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
@@ -192,7 +195,8 @@ class Adopt_adv(torch.optim.Optimizer):
                 # Initialize v_0 using NMF
                 _nnmf(vt_init, out=(state['mu_v_nmf'], state['mv_v_nmf']))
             else: # Fallback for non-factored tensors
-                state['exp_avg'] = torch.zeros_like(p, dtype=dtype) # m_0
+                if beta1 > 0:
+                    state['exp_avg'] = torch.zeros_like(p, dtype=dtype) # m_0
                 if self.use_AdEMAMix:
                     state['exp_avg_slow'] = torch.zeros_like(p, dtype=dtype)
                 state['exp_avg_sq'] = grad.square()   # v_0
@@ -202,7 +206,6 @@ class Adopt_adv(torch.optim.Optimizer):
             state['step'] += 1
             return
 
-        beta1, beta2 = group['betas']
         if self.use_AdEMAMix:
             beta3_ema = group['beta3_ema']
             alpha = group['alpha']
@@ -219,13 +222,14 @@ class Adopt_adv(torch.optim.Optimizer):
             d1, d2 = state['effective_shape']
 
             # Reconstruct m_{t-1}
-            mt = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
-            if not self.grams_moment:
-                if state['sign'].dtype != torch.uint8:
-                    state['sign'] = state['sign'].to(torch.uint8)
-                unpacked_sign = _unpack_bools(state['sign'], original_m=d2)
-                torch.where(unpacked_sign, mt, -mt, out=mt)
-                del unpacked_sign
+            if beta1 > 0:
+                mt = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
+                if not self.grams_moment:
+                    if state['sign'].dtype != torch.uint8:
+                        state['sign'] = state['sign'].to(torch.uint8)
+                    unpacked_sign = _unpack_bools(state['sign'], original_m=d2)
+                    torch.where(unpacked_sign, mt, -mt, out=mt)
+                    del unpacked_sign
 
             # Reconstruct AdEMAMix EMA
             if self.use_AdEMAMix:
@@ -253,25 +257,29 @@ class Adopt_adv(torch.optim.Optimizer):
             del denom
 
             # ADOPT Step B: Update momentum m_t using normalized gradient
-            if self.Simplified_AdEMAMix:
-                mt.mul_(beta1).add_(normalized_grad, alpha=1.0)
-            else:
-                mt.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
-            if self.grams_moment:
-                mt = grad_reshaped.sign() * mt.abs()
-            elif self.cautious_mask:
-                mask = (mt * grad_reshaped > 0).to(grad_reshaped.dtype)
-                mask.div_(mask.mean().clamp_(min=1e-3))
-                mt.mul_(mask)
-                del mask
+            if beta1 > 0:
+                if self.Simplified_AdEMAMix:
+                    mt.mul_(beta1).add_(normalized_grad, alpha=1.0)
+                else:
+                    mt.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
+                if self.grams_moment:
+                    mt = grad_reshaped.sign() * mt.abs()
+                elif self.cautious_mask:
+                    mask = (mt * grad_reshaped > 0).to(grad_reshaped.dtype)
+                    mask.div_(mask.mean().clamp_(min=1e-3))
+                    mt.mul_(mask)
+                    del mask
 
             if self.use_AdEMAMix:
                 mt_slow.mul_(beta3_ema).add_(normalized_grad, alpha=1.0 - beta3_ema)
-                update = torch.add(mt, mt_slow, alpha=alpha_t)
+                if beta1 > 0:
+                    update = torch.add(mt, mt_slow, alpha=alpha_t)
+                else:
+                    update = torch.add(normalized_grad, mt_slow, alpha=alpha_t)
             elif self.Simplified_AdEMAMix:
-                update = torch.add(mt, grad_reshaped, alpha=alpha_grad)
+                update = torch.add(mt, normalized_grad, alpha=alpha_grad)
             else:
-                update = mt.clone()
+                update = mt.clone() if beta1 > 0 else normalized_grad
 
             update = update.view(p.shape)
 
@@ -285,10 +293,11 @@ class Adopt_adv(torch.optim.Optimizer):
             del grad_reshaped
 
             # Compress and store new factors
-            if not self.grams_moment:
-                state['sign'] = _pack_bools(mt > 0)
-            _nnmf(mt.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
-            del mt
+            if beta1 > 0:
+                if not self.grams_moment:
+                    state['sign'] = _pack_bools(mt > 0)
+                _nnmf(mt.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
+                del mt
 
             if self.use_AdEMAMix:
                 state['sign_slow'] = _pack_bools(mt_slow > 0)
@@ -300,10 +309,7 @@ class Adopt_adv(torch.optim.Optimizer):
             del vt
 
         else: # Standard ADOPT logic for non-factored tensors
-            m, v = state['exp_avg'], state['exp_avg_sq'] # m_{t-1}, v_{t-1}
-            
-            if self.use_AdEMAMix:
-                m_slow = state['exp_avg_slow']
+            v = state['exp_avg_sq'] # v_{t-1}                
 
             # ADOPT Step A: Decorrelate g_t using v_{t-1}
             denom = v.sqrt()
@@ -318,10 +324,12 @@ class Adopt_adv(torch.optim.Optimizer):
             del denom
 
             # ADOPT Step B: Update momentum m_t
-            if self.Simplified_AdEMAMix:
-                m.mul_(beta1).add_(normalized_grad, alpha=1.0)
-            else:
-                m.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
+            if beta1 > 0:
+                m = state['exp_avg'] # m_{t-1},
+                if self.Simplified_AdEMAMix:
+                    m.mul_(beta1).add_(normalized_grad, alpha=1.0)
+                else:
+                    m.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
 
             if self.grams_moment:
                 m = grad.sign() * m.abs()
@@ -332,12 +340,16 @@ class Adopt_adv(torch.optim.Optimizer):
                 del mask
 
             if self.use_AdEMAMix:
+                m_slow = state['exp_avg_slow']
                 m_slow.mul_(beta3_ema).add_(normalized_grad, alpha=1.0 - beta3_ema)
-                update = torch.add(m, m_slow, alpha=alpha_t)
+                if beta1 > 0:
+                    update = torch.add(m, m_slow, alpha=alpha_t)
+                else:
+                    update = torch.add(normalized_grad, m_slow, alpha=alpha_t)
             elif self.Simplified_AdEMAMix:
-                update = torch.add(m, grad, alpha=alpha_grad)
+                update = torch.add(m, normalized_grad, alpha=alpha_grad)
             else:
-                update = m.clone()
+                update = m.clone() if beta1 > 0 else normalized_grad
 
             if self.use_atan2:
                 update.mul_(group['lr'] * 1.2732395447351628)

adv_optm/optim/Prodigy_adv.py

@@ -308,11 +308,11 @@ class Prodigy_adv(torch.optim.Optimizer):
                 if self.beta1 > 0:
                     update = torch.add(mt, mt_slow, alpha=alpha_t)
                 else:
-                    update = torch.add(grad_reshaped, mt_slow, alpha=alpha_t)
+                    update = torch.add(grad_reshaped.mul(self.d), mt_slow, alpha=alpha_t)
             elif self.Simplified_AdEMAMix:
                 update = torch.add(mt, grad_reshaped, alpha=alpha_grad * self.d)
             else:
-                update = mt.clone() if self.beta1 > 0 else grad_reshaped.clone()
+                update = mt.clone() if self.beta1 > 0 else grad_reshaped.mul(self.d)
             del grad_reshaped
 
             if group['use_atan2']:
@@ -362,11 +362,11 @@ class Prodigy_adv(torch.optim.Optimizer):
                 if self.beta1 > 0:
                     update = torch.add(exp_avg, exp_avg_slow, alpha=alpha_t)
                 else:
-                    update = torch.add(grad, exp_avg_slow, alpha=alpha_t)
+                    update = torch.add(grad.mul(self.d), exp_avg_slow, alpha=alpha_t)
             elif self.Simplified_AdEMAMix:
                 update = torch.add(exp_avg, grad, alpha=alpha_grad * self.d)
             else:
-                update = exp_avg.clone() if self.beta1 > 0 else grad.clone()
+                update = exp_avg.clone() if self.beta1 > 0 else grad.mul(self.d)
 
             exp_avg_sq.mul_(self.beta2).addcmul_(grad, grad.conj(), value=self.d * self.d * (1.0 - self.beta2))
 

{adv_optm-1.0.5.dist-info → adv_optm-1.0.6.dist-info}/METADATA

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

{adv_optm-1.0.5.dist-info → adv_optm-1.0.6.dist-info}/RECORD

@@ -1,9 +1,9 @@
-adv_optm/__init__.py,sha256=9sM4fP1pj9divFhLVUzHbBWe50H82H3FYIGVIHTHpkg,306
+adv_optm/__init__.py,sha256=dAbueuVEIGoYrYXx8UE4ATfFBH5wEKrpkXGPTjFH0r0,306
 adv_optm/optim/AdamW_adv.py,sha256=aTuYcJgd_EcZOrs6TDgBrBKw3wtU5LPzE5WvTBDDeEo,14317
-adv_optm/optim/Adopt_adv.py,sha256=KdEVSl2w1gRXFtz2fwCVT4i9inTspp-PQq3mobpa-9A,17476
+adv_optm/optim/Adopt_adv.py,sha256=FTpDDSlYruZDt1VVLgEI_bADiO8f26j-utQs7Gn2fFA,18108
 adv_optm/optim/Lion_Prodigy_adv.py,sha256=sGzhts9a6gHfCkuHTB5L9IrClo4c6UThzYYErBwqOaA,12844
 adv_optm/optim/Lion_adv.py,sha256=6G1CukJB_pC7l9HwFEuY1ydsNHZFabVmOvcHDsHHVuQ,8295
-adv_optm/optim/Prodigy_adv.py,sha256=8XUpu19BaBmHb-R9K3jgwySDbtVaLU1_Drtttc_zITs,22461
+adv_optm/optim/Prodigy_adv.py,sha256=G8xXLO9YBeLb9574uS0HpdY9w3ojblaV-PJFghUnToQ,22493
 adv_optm/optim/Simplified_AdEMAMix.py,sha256=tb3d6Cw_nGwcTzYUhDnKqyP7GzjD1hn8k4WqGG5lhmw,9813
 adv_optm/optim/__init__.py,sha256=pcP865H2j1tut2VfTUhzQh7V8TF_tzPjqFnjMfFed2k,382
 adv_optm/util/BF16_Stochastic_Rounding.py,sha256=Q5H0BcogmE4atP65dLoI21HKSf50lRdsBDfeF6v9Tbg,1548
@@ -12,8 +12,8 @@ adv_optm/util/NNMF.py,sha256=yRf5IP5Sjq0Uf0DxN0Q8NxEGSdD-f1ULziLVDOjY8K4,639
 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=qoyIF0jcLjs_vSEcsv36clw5LFNBEbifyXrrVxMH-G4,349
-adv_optm-1.0.5.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-1.0.5.dist-info/METADATA,sha256=ouxI4cwBQ2IPuOjrkA478XwSetGP6ku51vW1QxHIGcY,8422
-adv_optm-1.0.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-adv_optm-1.0.5.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-1.0.5.dist-info/RECORD,,
+adv_optm-1.0.6.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-1.0.6.dist-info/METADATA,sha256=3PslWXH0ysoiXU83vN3F9kWRw48fwUM4H1z1tMyEGvI,8422
+adv_optm-1.0.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+adv_optm-1.0.6.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-1.0.6.dist-info/RECORD,,