heavyball 0.21.1__tar.gz → 0.21.3__tar.gz

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

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

{heavyball-0.21.1 → heavyball-0.21.3}/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'], ea, p, lr, weight_decay)
+            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

{heavyball-0.21.1 → heavyball-0.21.3}/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'], ea, p, lr, weight_decay)
+            precond_grad_cached_(cached_q, ea, self.state_(p)['cache_expr'], p, lr, weight_decay)

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

@@ -11,6 +11,12 @@ 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):
+    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)
@@ -98,14 +104,14 @@ 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)
-            update_param_([p], self.clip_fn([new]), lr, weight_decay)
+            _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],

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

@@ -968,7 +968,7 @@ 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):
     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)
+    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)
 
 

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

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

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

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