heavyball 0.21.4__tar.gz → 0.21.6__tar.gz

{heavyball-0.21.4 → heavyball-0.21.6}/PKG-INFO

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

{heavyball-0.21.4 → heavyball-0.21.6}/heavyball/delayed_psgd.py

@@ -10,9 +10,10 @@ 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_deca, clip_fn):
+# TODO: E1123 00:51:55.423000 159394 site-packages/torch/_guards.py:283] [5/0] Error while creating guard:
+# E1123 00:51:55.423000 159394 site-packages/torch/_guards.py:283] [5/0] Name: "G['psgd_precond_grad'].__defaults__[0]"
+@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], clip_fn([new]), lr, weight_decay)
 

{heavyball-0.21.4 → heavyball-0.21.6}/heavyball/utils.py

@@ -965,7 +965,7 @@ class PSGDBase(StatefulOptimizer):
                         psgd_balance_Q(q)
 
 
-@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
 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)

{heavyball-0.21.4 → heavyball-0.21.6}/heavyball.egg-info/PKG-INFO

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

{heavyball-0.21.4 → heavyball-0.21.6}/setup.py

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

{heavyball-0.21.4 → heavyball-0.21.6}/test/test_bf16_params.py

@@ -20,7 +20,7 @@ def get_memory():
 
 @pytest.mark.parametrize("opt", heavyball.__all__)
 @pytest.mark.parametrize("size,depth", [(256, 2)])
-def test_foreach(opt, size, depth: int, iterations: int = 16, outer_iterations: int = 3):
+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')