heavyball 0.21.0__py3-none-any.whl → 0.21.2__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'], 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'], p, lr, weight_decay)

heavyball/delayed_psgd.py

@@ -5,12 +5,18 @@ 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
 
 
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_psgd_precond_grad_(q, exprs, ea, p, lr, weight_decay):
+    new = psgd_precond_grad(q, exprs, ea)
+    update_param_([p], self.clip_fn([new]), lr, weight_decay)
+
+
 class ForeachDelayedPSGD(PSGDBase):
     """
     Implements PSGD with off-by-one preconditioning (akin to ADOPT and SOAP)
@@ -39,7 +45,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:
@@ -98,15 +104,15 @@ class ForeachDelayedPSGD(PSGDBase):
         stochastic_lerp_(exp_avg_list, grad_list, beta_debias(beta, group["step"]))
 
         lr = -warmup(lr, group['step'], group['warmup_steps'])
+        lr = torch.empty((), dtype=torch.float32, device=grad_list[0].device).fill_(lr)
 
         Q_list, exp_avg_list = list(Q_list), list(exp_avg_list)
         for i, (p, g) in enumerate(zip(p_list, grad_list)):
             q_orig = Q_list.pop(0)
             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)
+            _compilable_psgd_precond_grad_(q, state["exprs"], ea, p, 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

@@ -965,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.21.0.dist-info → heavyball-0.21.2.dist-info}/METADATA

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

{heavyball-0.21.0.dist-info → heavyball-0.21.2.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=jBMA5MI4EpOIKww0iQ0zkRO4g1EPPvWHpNWS5ndgDo8,6660
+heavyball/cached_psgd_kron.py,sha256=5qENxAs2N0YRWI0CHFK30bdffxVjOagfAonODKnkDnI,6578
+heavyball/delayed_psgd.py,sha256=b0Fd-3IoqC5Q0eSCGUGkOmKOPHsQke_Ozv4uWVEAZoU,5929
 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=H8RsADNAXVbjQ9wWstYIKkXMq9E81aUF1j-2wfCeSLA,36471
-heavyball-0.21.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.21.0.dist-info/METADATA,sha256=hbXhr4XcPAkgfW8hpgoPRPrUoKeTCTvhZdofj4h8_8c,11926
-heavyball-0.21.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.21.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.21.0.dist-info/RECORD,,
+heavyball/utils.py,sha256=9PaYhieWXIo6uE-Bb48N45wOg-AmQEkmNnTXjtKByfo,37211
+heavyball-0.21.2.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.21.2.dist-info/METADATA,sha256=XvA88WVYQdXXj04iYuPDUWM2rKiUZInJkNvC5BM-zqA,11926
+heavyball-0.21.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.21.2.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.21.2.dist-info/RECORD,,