heavyball 0.17.1__py3-none-any.whl → 0.17.3__py3-none-any.whl

heavyball/utils.py

@@ -332,6 +332,16 @@ def promote(x):
     return x
 
 
+def min_dtype(xs: List[torch.Tensor]):
+    dtypes = [x.dtype for x in xs]
+    for d in (torch.float32, torch.bfloat16, torch.float16):
+        if all(d == x for x in dtypes):
+            return d
+        if all(d in (x, torch.float32, torch.float64) for x in dtypes):
+            return d
+    return torch.float32
+
+
 def update_preconditioner(grad, state, max_precond_dim, precondition_1d, beta, update_precond):
     """
     Updates the preconditioner matrices and the eigenbases (L, R, Q_L, Q_R in the paper).
@@ -471,13 +481,8 @@ def copy_stochastic_list_(target: List[torch.Tensor], source: List[torch.Tensor]
         copy_stochastic_(t, s)
 
 
-def copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
-    if target.data_ptr() == source.data_ptr():
-        return
-    if target.dtype != torch.bfloat16:
-        set_(target, source)
-        return
-
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+def _compilable_copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
     """Taken as-is from https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905"""
     # create a random 16 bit integer
     result = torch.randint_like(source, dtype=torch.int32, low=0, high=(1 << 16))
@@ -492,6 +497,15 @@ def copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
     target.copy_(result.view(dtype=torch.float32))
 
 
+def copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
+    if target.data_ptr() == source.data_ptr():
+        return
+    if target.dtype != torch.bfloat16 or source.dtype not in (torch.float16, torch.float32, torch.float64):
+        set_(target, source)
+        return
+    _compilable_copy_stochastic_(target, source)
+
+
 def update_param_(param: List[torch.Tensor], update: List[torch.Tensor], lr: float, decay: float,
                   add_fn: callable = None):
     param32 = [promote(p) for p in param]
@@ -602,7 +616,8 @@ def psgd_balance_Q(Q_in):
 
 
 def psgd_calc_A_and_conjB(exprA, G, Q, V):
-    A = torch.einsum(exprA, *Q, G)
+    md = min_dtype(Q)
+    A = torch.einsum(exprA, *[q.to(md) for q in Q], G.to(md))
     order = G.dim()
     p = list(range(order))
     conjB = torch.permute(V.conj(), p[1:] + p[:1])
@@ -653,7 +668,10 @@ def psgd_update_precond(Q, exprs, V, G, step, tiny):
 
         term2 += term1  # a + b
         term1 *= 2  # 2a
-        term1 -= term2  # 2a - (a + b) == a - b
+        if term1.dtype == term2.dtype:
+            term1 -= term2  # 2a - (a + b) == a - b
+        else:
+            term1 = term1 - term2
 
         term1 *= step
         norm = term2.norm(float('inf'))
@@ -669,7 +687,8 @@ def psgd_update_precond(Q, exprs, V, G, step, tiny):
 @decorator
 def psgd_precond_grad(Q, exprs, G, inplace: bool = False):
     """Precondition gradient G with preconditioner Q."""
-    out = torch.einsum(exprs[-1], *[q.conj() for q in Q], *Q, G.to(Q[0].dtype))
+    md = min_dtype(Q)
+    out = torch.einsum(exprs[-1], *[q.conj().to(md) for q in Q], *[q.to(md) for q in Q], G.to(md))
     if inplace:
         set_(G, out)
         return G
@@ -787,14 +806,15 @@ class PSGDBase(StatefulOptimizer):
             if g.dim() > 1:
                 psgd_balance_Q(q)
 
-    def do_update(self, p_list, grad_list, q_list, precond_lr, original_q: Optional[List] = None, store_triu_as_line=False):
+    def do_update(self, p_list, grad_list, q_list, precond_lr, original_q: Optional[List] = None,
+                  store_triu_as_line=False):
         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:
                 if store_triu_as_line:
                     update_triu_(original_q[i], Q)
                 else:
-                    copy_stochastic_(original_q[i], Q)
+                    copy_stochastic_list_(original_q[i], Q)
 
 
 def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.001, flat_start=250):

{heavyball-0.17.1.dist-info → heavyball-0.17.3.dist-info}/METADATA

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

{heavyball-0.17.1.dist-info → heavyball-0.17.3.dist-info}/RECORD

@@ -16,9 +16,9 @@ heavyball/precond_schedule_sfpsoap.py,sha256=vq7jd302refKPa_9X2lkOTOtCCcTBVByPdo
 heavyball/psgd_kron.py,sha256=2IpPj2TOExNGm8hSewi3er2GczJRNgC7r2J5yYSSA_0,5998
 heavyball/pure_psgd.py,sha256=uA7W9a3Qm1sxHQhtNxaUYrmE5x55lP5iJOKy_qT8XaQ,5341
 heavyball/schedule_free_palm_foreach_soap.py,sha256=zkcikH5wWbzq4kOrmBjilvY3iWzuUddcv2HNEPKr3MI,6366
-heavyball/utils.py,sha256=Jqh7VdWGeiSdwaPtUNB9l14wuuFPSReLaTwJA3juFbM,28765
-heavyball-0.17.1.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.17.1.dist-info/METADATA,sha256=2FAgCpyuH4G-B_m0mhbl-sdkMizS1sd8oNmNkPpAKN0,11810
-heavyball-0.17.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.17.1.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.17.1.dist-info/RECORD,,
+heavyball/utils.py,sha256=wEwg-TElm56CEf-Q_F7wgJdz17P4Hp0CakALGJj1090,29563
+heavyball-0.17.3.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.17.3.dist-info/METADATA,sha256=QMYWssjAbuY_eYSGRVuSOVeFrMiq7YZA4-qJe8E8u74,11810
+heavyball-0.17.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.17.3.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.17.3.dist-info/RECORD,,