heavyball 0.21.3__tar.gz → 0.21.5__tar.gz

{heavyball-0.21.3 → heavyball-0.21.5}/PKG-INFO

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

{heavyball-0.21.3 → heavyball-0.21.5}/heavyball/cached_delayed_psgd_kron.py

@@ -120,7 +120,7 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
             q_orig = Q_list.pop(0)
             ea = exp_avg_list.pop(0)
 
-            precond_grad_cached_(cached_q, ea, self.state_(p)['cache_expr'], p, lr, weight_decay)
+            precond_grad_cached_(cached_q, ea, self.state_(p)['cache_expr'], p, lr, weight_decay, self.clip_fn)
 
             if should_update:
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig

{heavyball-0.21.3 → heavyball-0.21.5}/heavyball/cached_psgd_kron.py

@@ -128,4 +128,4 @@ class ForeachCachedPSGDKron(PSGDBase):
                     else:
                         torch.mul(q_.conj(), q_, out=c_)
 
-            precond_grad_cached_(cached_q, ea, self.state_(p)['cache_expr'], p, lr, weight_decay)
+            precond_grad_cached_(cached_q, ea, self.state_(p)['cache_expr'], p, lr, weight_decay, self.clip_fn)

{heavyball-0.21.3 → heavyball-0.21.5}/heavyball/delayed_psgd.py

@@ -11,10 +11,10 @@ from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust
     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):
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+def _compilable_psgd_precond_grad_(q, exprs, ea, p, lr, weight_decay, clip_fn):
     new = psgd_precond_grad(q, exprs, ea)
-    update_param_([p], self.clip_fn([new]), lr, weight_decay)
+    update_param_([p], clip_fn([new]), lr, weight_decay)
 
 
 class ForeachDelayedPSGD(PSGDBase):
@@ -62,7 +62,7 @@ class ForeachDelayedPSGD(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)
+                        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):
@@ -111,7 +111,7 @@ class ForeachDelayedPSGD(PSGDBase):
             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
-            _compilable_psgd_precond_grad_(q, state["exprs"], ea, p, lr, weight_decay)
+            _compilable_psgd_precond_grad_(q, self.state_(p)["exprs"], ea, p, lr, weight_decay, self.clip_fn)
             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],

{heavyball-0.21.3 → heavyball-0.21.5}/heavyball/p_adam.py

@@ -6,7 +6,7 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 
 import torch
 
-from heavyball.utils import triu_to_line, line_to_triu, identity
+from heavyball.utils import triu_to_line, line_to_triu, identity, stochastic_lerp_
 from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, PSGDBase, exp_avg_sq_, beta_debias, \
     split_p_and_g_in_group, promote
 
@@ -100,7 +100,7 @@ class ForeachPaLMPAdam(PSGDBase):
             for g, p, q_, q_orig in zip(grad_list, p_list, Q_triu, Q_list):
                 q32 = [promote(qq_) for qq_ in q_]
                 self.do_update(group, [p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
-        torch._foreach_lerp_(exp_avg, grad_list, 1 - beta_debias(group['beta'], group['step']))
+        stochastic_lerp_(exp_avg, grad_list, 1 - beta_debias(group['beta'], group['step']))
 
         beta2 = 1 - group['step'] ** -group['beta2_scale']
 

{heavyball-0.21.3 → heavyball-0.21.5}/heavyball/utils.py

@@ -966,17 +966,17 @@ class PSGDBase(StatefulOptimizer):
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
-def _compilable_precond_grad_cached_(cached_q, ea, expr, param, lr, weight_decay):
+def _compilable_precond_grad_cached_(cached_q, ea, expr, param, lr, weight_decay, clip_fn):
     md = min_dtype(cached_q + [ea])
     new = torch.einsum(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)
+    update_param_([param], 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):
+                         weight_decay: float, clip_fn):
     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)
+    _compilable_precond_grad_cached_(cached_q, ea, expr, param, lr, weight_decay, clip_fn)
 
 
 def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.001, flat_start=250):

{heavyball-0.21.3 → heavyball-0.21.5}/heavyball.egg-info/PKG-INFO

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

{heavyball-0.21.3 → heavyball-0.21.5}/setup.py

@@ -10,7 +10,7 @@ setuptools.setup(
     name='heavyball',
     license='BSD',
     description='Efficient optimizers',
-    version='0.21.3',
+    version='0.21.5',
     long_description=README,
     url='https://github.com/clashluke/heavyball',
     packages=setuptools.find_packages(),

{heavyball-0.21.3 → heavyball-0.21.5}/test/test_bf16_params.py

@@ -22,6 +22,8 @@ def get_memory():
 @pytest.mark.parametrize("size,depth", [(256, 2)])
 def test_foreach(opt, size, depth: int, iterations: int = 512, outer_iterations: int = 3):
     set_torch()
+    if 'psgd' not in opt.lower() and 'padam' not in opt.lower():
+        raise pytest.skip('Only PSGD and PaLMPAdam are supported')
     opt = getattr(heavyball, opt)