heavyball 0.18.1__py3-none-any.whl → 0.18.2__py3-none-any.whl

heavyball/cached_delayed_psgd_kron.py

@@ -113,6 +113,8 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
             q_orig = Q_list.pop(0)
             ea = exp_avg_list.pop(0)
 
+            new = torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea)
+
             if self.should_update(group):
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig
                 q32 = [promote(q_) for q_ in q]
@@ -124,7 +126,7 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                     else:
                         torch.mul(q_.conj(), q_, out=c_)
 
-            set_(g, torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea))
+            set_(g, new)
         grad_list = self.clip_fn(grad_list)
 
         lr = -warmup(lr, group['step'], group['warmup_steps'])

heavyball/cached_psgd_kron.py

@@ -111,8 +111,6 @@ class ForeachCachedPSGDKron(PSGDBase):
             q_orig = Q_list.pop(0)
             ea = exp_avg_list.pop(0)
 
-            new = torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea)
-
             if self.should_update(group):
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig
                 q32 = [promote(q_) for q_ in q]
@@ -124,7 +122,7 @@ class ForeachCachedPSGDKron(PSGDBase):
                     else:
                         torch.mul(q_.conj(), q_, out=c_)
 
-            set_(g, new)
+            set_(g, torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea))
 
         grad_list = self.clip_fn(grad_list)
 

heavyball/utils.py

@@ -817,14 +817,15 @@ class PSGDBase(StatefulOptimizer):
 
     def do_update(self, group, p_list, grad_list, q_list, precond_lr, original_q: Optional[List] = None,
                   store_triu_as_line=False):
-        if self.should_update(group, self.balance_probability, 'balance_prob'):
-            for g, q in zip(grad_list, q_list):
-                if g.dim() > 1:
-                    psgd_balance_Q(q)
-
         for i, (p, grad, Q) in enumerate(zip(p_list, grad_list, q_list)):
             psgd_update_precond(Q, self.state_(p)["exprs"], torch.randn_like(grad), grad, precond_lr, self._tiny)
-            if original_q:
+
+        for g, q in zip(grad_list, q_list):
+            if g.dim() > 1:
+                psgd_balance_Q(q)
+
+        if original_q:
+            for q in q_list:
                 if store_triu_as_line:
                     update_triu_(original_q[i], Q)
                 else:

{heavyball-0.18.1.dist-info → heavyball-0.18.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.18.1
+Version: 0.18.2
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler

{heavyball-0.18.1.dist-info → heavyball-0.18.2.dist-info}/RECORD

@@ -1,6 +1,6 @@
 heavyball/__init__.py,sha256=iqP428JWwwx-XDOZ0nUdbCkOLEyfoqVyWZLQLAcwxaw,2214
-heavyball/cached_delayed_psgd_kron.py,sha256=IXGaMBqHidw7FikWjpc1yiuCLZrbCLi45rRvL2OfKxU,6399
-heavyball/cached_psgd_kron.py,sha256=rxrVaXYoivdAME6R8xoWECsIZ1DL8mfAw0YSu94pCj4,6402
+heavyball/cached_delayed_psgd_kron.py,sha256=t8XXsl91lINY0iB5Cn5aQDyjLxN2itGiA97Pur4mEkY,6422
+heavyball/cached_psgd_kron.py,sha256=NnzgJB11xPfi5NrHl3OsQkgs-fxeR_tsHdfXeDiXxbE,6379
 heavyball/delayed_psgd.py,sha256=ylLNHglvjnkYAmJwcl1TtPA4PXKPaOv1YHVt0JVabMA,5551
 heavyball/foreach_adamw.py,sha256=h_ar0ZZRM0q_wxuEkxEOEYe-2p-mB4OMgAHivrUnPl8,1777
 heavyball/foreach_adopt.py,sha256=ogOw2JjwEQNj7AKlweAphQFdMJ_GcMDm-RyDvEzugoc,1911
@@ -16,9 +16,9 @@ heavyball/precond_schedule_sfpsoap.py,sha256=vq7jd302refKPa_9X2lkOTOtCCcTBVByPdo
 heavyball/psgd_kron.py,sha256=KhZnV5MpigAEfJfvYI7ApF1GQ8ZWWXl7g5nYueWKYDQ,5438
 heavyball/pure_psgd.py,sha256=qPQ46pp7DWyQ1afBin2bqFVhaRhjt7RjXm6VuM_2sxg,4851
 heavyball/schedule_free_palm_foreach_soap.py,sha256=zkcikH5wWbzq4kOrmBjilvY3iWzuUddcv2HNEPKr3MI,6366
-heavyball/utils.py,sha256=sojDNo94l-jAPDSnhEM5EI6D83SHG_hTnJOkdnDquSI,30492
-heavyball-0.18.1.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.18.1.dist-info/METADATA,sha256=PeclJ6CUn0l-ycvQoW4uN-BTUml3xXdiPitBMidjryA,11810
-heavyball-0.18.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.18.1.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.18.1.dist-info/RECORD,,
+heavyball/utils.py,sha256=LTkqF5hP6z9De4_jDLF0HQYUwz1MkJEOTdMOYyH5D0k,30426
+heavyball-0.18.2.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.18.2.dist-info/METADATA,sha256=H074F1jIGrFvWdTdeRmhGXoxJ5V0zH7lIEhZ-LSP3Mc,11810
+heavyball-0.18.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.18.2.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.18.2.dist-info/RECORD,,