heavyball 0.20.1__py3-none-any.whl → 0.21.1__py3-none-any.whl

heavyball/cached_delayed_psgd_kron.py

@@ -7,9 +7,10 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 from typing import Optional
 
 import torch
+from heavyball.utils import min_dtype, precond_grad_cached_
 
 from .utils import update_param_, warmup, init_Q_exprs, trust_region_clip_, PSGDBase, split_p_and_g_in_group, \
-    line_to_triu, triu_to_line, einsum_base, promote, stochastic_lerp_, beta_debias
+    line_to_triu, triu_to_line, einsum_base, promote, stochastic_lerp_, beta_debias, precond_grad_cached_
 
 
 class ForeachCachedDelayedPSGDKron(PSGDBase):
@@ -59,7 +60,8 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
                         momentum_into_precond_update=momentum_into_precond_update, precond_lr=precond_lr,
                         precond_init_scale=precond_init_scale, step=0, warmup_steps=warmup_steps, merge_dims=merge_dims,
-                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype, storage_dtype=storage_dtype)
+                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype,
+                        storage_dtype=storage_dtype)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
@@ -118,7 +120,7 @@ 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)
+            precond_grad_cached_(cached_q, ea, self.state_(p)['cache_expr'], ea, p, lr, weight_decay)
 
             if should_update:
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig
@@ -130,5 +132,3 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                         torch.matmul(q_.T.conj(), q_, out=c_)
                     else:
                         torch.mul(q_.conj(), q_, out=c_)
-
-            update_param_([p], self.clip_fn([new]), lr, weight_decay)

heavyball/cached_psgd_kron.py

@@ -9,7 +9,7 @@ from typing import Optional
 import torch
 
 from .utils import update_param_, warmup, init_Q_exprs, trust_region_clip_, PSGDBase, split_p_and_g_in_group, \
-    line_to_triu, triu_to_line, einsum_base, promote, stochastic_lerp_, beta_debias
+    line_to_triu, triu_to_line, einsum_base, promote, stochastic_lerp_, beta_debias, precond_grad_cached_
 
 
 class ForeachCachedPSGDKron(PSGDBase):
@@ -128,5 +128,4 @@ class ForeachCachedPSGDKron(PSGDBase):
                     else:
                         torch.mul(q_.conj(), q_, out=c_)
 
-            g = torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea)
-            update_param_([p], self.clip_fn([g]), lr, weight_decay)
+            precond_grad_cached_(cached_q, ea, self.state_(p)['cache_expr'], ea, p, lr, weight_decay)

heavyball/delayed_psgd.py

@@ -5,8 +5,8 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 """
 
 import torch
-
 from heavyball.utils import stochastic_lerp_, beta_debias
+
 from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust_region_clip_, PSGDBase, \
     split_p_and_g_in_group, triu_to_line, line_to_triu, promote
 
@@ -39,7 +39,7 @@ class ForeachDelayedPSGD(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: callable = None, store_triu_as_line: bool = True, foreach: bool = True,
-                 q_dtype='float32', stochastic_schedule: bool = True, storage_dtype:str='float32', #
+                 q_dtype='float32', stochastic_schedule: bool = True, storage_dtype: str = 'float32',  #
                  # expert parameters
                  precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
@@ -105,8 +105,8 @@ class ForeachDelayedPSGD(PSGDBase):
             ea = exp_avg_list.pop(0)
             q = line_to_triu(q_orig) if store_triu_as_line else q_orig
             new = psgd_precond_grad(q, self.state_(p)["exprs"], ea)
+            update_param_([p], self.clip_fn([new]), lr, weight_decay)
             if should_update:
                 q32 = [promote(q_) for q_ in q]
                 self.do_update(group, [p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig],
                                store_triu_as_line)
-            update_param_([p], self.clip_fn([new]), lr, weight_decay)

heavyball/utils.py

@@ -329,23 +329,33 @@ def get_orthogonal_matrix(mat):
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
-def stochastic_lerp_(x: List[torch.Tensor], y: List[torch.Tensor], a: Union[float, int, torch.Tensor]):
-    x32 = [promote(x_) for x_ in x]
-    y32 = [promote(y_) for y_ in y]
+def _compilable_stochastic_lerp_(x: List[torch.Tensor], y: List[torch.Tensor], a: Union[float, int, torch.Tensor]):
+    for x_, y_ in zip(x, y):
+        x32 = promote(x_)
+        y32 = promote(y_)
+        x32.lerp_(y32, a)
+        copy_stochastic_(x_, x32)
 
-    torch._foreach_lerp_(x32, y32, a)
 
-    copy_stochastic_list_(x, x32)
+def stochastic_lerp_(x: List[torch.Tensor], y: List[torch.Tensor], a: Union[float, int, torch.Tensor]):
+    if not isinstance(a, torch.Tensor):
+        a = torch.empty((), dtype=torch.float32, device=x[0].device).fill_(a)
+    _compilable_stochastic_lerp_(x, y, a)
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
-def stochastic_add_(x: List[torch.Tensor], y: List[torch.Tensor], alpha: Union[float, int, torch.Tensor]):
-    x32 = [promote(x_) for x_ in x]
-    y32 = [promote(y_) for y_ in y]
+def _compilable_stochastic_add_(x: List[torch.Tensor], y: List[torch.Tensor], alpha: Union[float, int, torch.Tensor]):
+    for x_, y_ in zip(x, y):
+        x32 = promote(x_)
+        y32 = promote(y_)
+        x32.add_(y32, alpha=alpha)
+        copy_stochastic_(x_, x32)
 
-    [x_.add_(y_, alpha=alpha) for x_, y_ in zip(x32, y32)]
 
-    copy_stochastic_list_(x, x32)
+def stochastic_add_(x: List[torch.Tensor], y: List[torch.Tensor], alpha: Union[float, int, torch.Tensor]):
+    if not isinstance(alpha, torch.Tensor):
+        alpha = torch.empty((), dtype=torch.float32, device=x[0].device).fill_(alpha)
+    _compilable_stochastic_add_(x, y, alpha)
 
 
 @decorator
@@ -572,7 +582,7 @@ def copy_stochastic_list_(target: List[torch.Tensor], source: List[torch.Tensor]
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
-def exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step):
+def _compilable_exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -585,6 +595,18 @@ def exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step):
     return denom
 
 
+def exp_avg_(exp_avg: List[torch.Tensor], exp_avg_sq: List[torch.Tensor], grad: List[torch.Tensor],
+             grad_projected: List[torch.Tensor], beta1: float, beta2: float, step: int):
+    if isinstance(beta1, float):
+        beta1 = torch.empty((), dtype=torch.float32, device=exp_avg[0].device).fill_(beta1)
+    if isinstance(beta2, float):
+        beta2 = torch.empty((), dtype=torch.float32, device=exp_avg[0].device).fill_(beta2)
+    if isinstance(step, int):
+        step = torch.empty((), dtype=torch.int32, device=exp_avg[0].device).fill_(step)
+    denom = _compilable_exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step)
+    return denom
+
+
 # this can be dynamic for most optimizers - just not for PSGD. So, it's disabled for all
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True)
 def _compilable_copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
@@ -734,7 +756,7 @@ def psgd_calc_A_and_conjB(exprA, G, Q):
     A = torch.einsum(exprA, *[q.to(md) for q in Q], G.to(md)).to(G.dtype)
     order = G.dim()
     p = list(range(order))
-    conjB = torch.randn_like(G.shape[1:] + G.shape[:1], dtype=promote(G.dtype), device=G.device)
+    conjB = torch.randn(G.shape[1:] + G.shape[:1], dtype=promote(G.dtype), device=G.device)
     Q = [promote(q) for q in Q]
     for i, q in enumerate(Q):
         if q.dim() <= 1:
@@ -756,17 +778,17 @@ def psgd_lb(A, max_abs):
     a0 = torch.einsum('ij,ij->j', A, A)
     i = torch.argmax(a0)
 
-    x = torch.index_select(a, 1, i).flatten().contiguous()
+    x = torch.index_select(A, 1, i).flatten().contiguous()
 
-    x = torch.einsum('i,ij->j', x_, a)
+    x = torch.einsum('i,ij->j', x, A)
     x /= x.norm()
-    x = torch.einsum('j,kj->k', x_, a)
+    x = torch.einsum('j,kj->k', x, A)
     x = x.norm()
     x *= max_abs
     return x
 
 
-@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
 def psgd_update_precond(Q, exprs, G, precond_lr, tiny, oq, store_triu_as_line):
     """Update Kronecker product preconditioner Q with pair (V, G)."""
     exprA, exprGs, _ = exprs
@@ -799,7 +821,7 @@ def psgd_update_precond(Q, exprs, G, precond_lr, tiny, oq, store_triu_as_line):
         stochastic_add_([o], [term1], -1)
 
 
-@decorator
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
 def psgd_precond_grad(Q, exprs, G, inplace: bool = False):
     """Precondition gradient G with preconditioner Q."""
     md = min_dtype(Q)
@@ -943,6 +965,20 @@ class PSGDBase(StatefulOptimizer):
                         psgd_balance_Q(q)
 
 
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_precond_grad_cached_(cached_q, ea, expr, param, lr, weight_decay):
+    md = min_dtype(cached_q + [ea])
+    new = torch.einsum(self.state_(p)['cache_expr'], *[c_.to(md) for c_ in cached_q], ea.to(md)).to(torch.float32)
+    update_param_([param], self.clip_fn([new]), lr, weight_decay)
+
+
+def precond_grad_cached_(cached_q: List[torch.Tensor], ea: torch.Tensor, expr: str, param: torch.Tensor, lr: float,
+                         weight_decay: float):
+    if isinstance(lr, float):
+        lr = torch.empty((), dtype=torch.float32, device=param.device).fill_(lr)
+    _compilable_precond_grad_cached_(cached_q, ea, expr, param, lr, weight_decay)
+
+
 def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.001, flat_start=250):
     """Anneal preconditioner update probability during beginning of training.
 

{heavyball-0.20.1.dist-info → heavyball-0.21.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.20.1
+Version: 0.21.1
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler
@@ -32,7 +32,7 @@ A simple package of efficient optimizers
 The goal is not to thrive for completeness, full maintenance or abstraction, but instead to provide a simple
 largely static alternative to `torch.optim` with more and better optimizers.
 
-Currently (2024-11-22, 0.19.1), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
+Currently (2024-11-22, 0.21.0), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
 recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psgd_efficiency.md)).
 
 ## Features

{heavyball-0.20.1.dist-info → heavyball-0.21.1.dist-info}/RECORD

@@ -1,7 +1,7 @@
 heavyball/__init__.py,sha256=iqP428JWwwx-XDOZ0nUdbCkOLEyfoqVyWZLQLAcwxaw,2214
-heavyball/cached_delayed_psgd_kron.py,sha256=apVzESMaQ8uxunHvfvYfyWA8HLbS25wQSd3j_YNEjGs,6603
-heavyball/cached_psgd_kron.py,sha256=3IETfsC0Ufu_8TPfo9SByGmztwjW6ktSFPwHNrUWkys,6601
-heavyball/delayed_psgd.py,sha256=0LaazbiBZOdx78EDS-945cW3bmeORjUvdFOGqdw3aMs,5631
+heavyball/cached_delayed_psgd_kron.py,sha256=QMSQ1es2ViAf6KWEuePihyjqWfTN22-RNYve31CCZdI,6664
+heavyball/cached_psgd_kron.py,sha256=NIWDUUrocU8rNR5HsIyR28XtxZ3OX7uKa25-J21YCJk,6582
+heavyball/delayed_psgd.py,sha256=I1EhQU7CUBLY94fXIdrzLtgIpfnQHU5XWaKf0S2Ls34,5635
 heavyball/foreach_adamw.py,sha256=Rb5U80cgUcEqlEbUU250UTWdoqA7nyiqkV5w1U4bWX4,2445
 heavyball/foreach_adopt.py,sha256=ecdi1fKg9i087OGjtKWVbE_DD6Yf4pvpzv4ELCcusvQ,3211
 heavyball/foreach_laprop.py,sha256=vi6C_gfjXxw5uN0KHgzxI9itUI1dcgOf3ufoO_VVMp0,2471
@@ -16,9 +16,9 @@ heavyball/precond_schedule_sfpsoap.py,sha256=PNneiOkrRyV1yIZn91lPmYofd1_OiLqJTDy
 heavyball/psgd_kron.py,sha256=RSLJi5FnSmYjvYBufNDClnYRm-eN_Kpa1Ar2tNP6-X0,5824
 heavyball/pure_psgd.py,sha256=LZK0qmvZkBF8g00evaVLtW-sIUJmdoxag1K7O26AqEo,4820
 heavyball/schedule_free_palm_foreach_soap.py,sha256=_AqrnChY6iDlQUkF2YUxS7eLjSWCIuvEUOHvMHVM1yY,6873
-heavyball/utils.py,sha256=14vt4r_MeTsp1q3m0lpgF-Q3PCJg6GLGJrhjRxnbWwQ,35174
-heavyball-0.20.1.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.20.1.dist-info/METADATA,sha256=qzF2P7e2EREeTy_4h85tvUY53omjNm32z83CUHTqt3U,11926
-heavyball-0.20.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.20.1.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.20.1.dist-info/RECORD,,
+heavyball/utils.py,sha256=9PaYhieWXIo6uE-Bb48N45wOg-AmQEkmNnTXjtKByfo,37211
+heavyball-0.21.1.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.21.1.dist-info/METADATA,sha256=NzvyCgg7cX4R_ouN9uQ1wE4KC6O2aQGwZLSRYMCSR7Q,11926
+heavyball-0.21.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.21.1.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.21.1.dist-info/RECORD,,