heavyball 1.1.0__tar.gz → 1.1.1__tar.gz

{heavyball-1.1.0 → heavyball-1.1.1}/PKG-INFO

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

{heavyball-1.1.0 → heavyball-1.1.1}/heavyball/chainable.py

@@ -140,16 +140,14 @@ class SkipUpdate(ValueError):
 @zero_guard("exp_avg")
 @no_state
 def exp_avg(group, update, grad, param, exp_avg):
-    utils.stochastic_lerp_(exp_avg, update, utils.beta_debias(utils.get_beta1(group), group["step"]))
-    return exp_avg
+    return utils.scale_by_exp_avg_(exp_avg, update, utils.beta_debias(utils.get_beta1(group), group["step"]))
 
 
 @zero_guard("exp_avg_sq")
 @no_state
 def scale_by_exp_avg_sq(group, update, grad, param, exp_avg_sq):
-    out = utils.scale_by_exp_avg_sq_(exp_avg_sq, update, utils.beta_debias(utils.get_beta2(group), group["step"]),
-                                     group['eps'])
-    return out
+    return utils.scale_by_exp_avg_sq_(exp_avg_sq, update, utils.beta_debias(utils.get_beta2(group), group["step"]),
+                                      group['eps'])
 
 
 @zero_guard("exp_avg", "exp_avg_sq")
@@ -350,9 +348,12 @@ def _fused_cached_psgd_precond_grad(group, grad, param, cached, cache_expr, expr
 @no_state_no_foreach
 def scale_by_psgd(group, update, grad, param, Q, exprs, Q_cache, cache_expr: str, cached: bool = False,
                   prob: Optional[callable] = None):
+    old = update
+    update = update.to(memory_format=torch.contiguous_format)
     Q_mat = utils.line_to_triu(Q) if group['store_triu_as_line'] else Q
     _update_psgd_precond(group, param, update, Q_mat, Q, exprs, prob)
-    return _cached_psgd_precond_grad(False, cache_expr, exprs, update, Q_mat, Q_cache)
+    out = _cached_psgd_precond_grad(False, cache_expr, exprs, update, Q_mat, Q_cache)
+    return torch.as_strided(out, old.shape, old.stride())
 
 
 @general_guard("Q", "exprs", ("Q_cache", None), ("cache_expr", None), init_fn=_init_psgd)
@@ -400,7 +401,7 @@ def apply_to_idx(fn, idx):
 
 
 def chain(state: Union[callable, dict], group, grad, param, *fns):
-    update = [torch.clone(g) for g in grad]
+    update = [torch.clone(g, memory_format=torch.preserve_format) for g in grad]
     skip_update = False
     for fn in fns:
         try:

{heavyball-1.1.0 → heavyball-1.1.1}/heavyball/utils.py

@@ -201,6 +201,21 @@ def scale_by_exp_avg_sq_(exp_avg_sq, grad, beta2, eps):
     return grad
 
 
+@decorator_knowngood
+def _compilable_exp_avg_(state, grad, beta):
+    s32, g32 = [list(map(promote, x)) for x in (state, grad)]
+    [s.lerp_(g, beta) for s, g in zip(s32, g32)]
+    copy_stochastic_list_(state, s32)
+    copy_stochastic_list_(grad, s32)
+
+
+def scale_by_exp_avg_(state, grad, beta):
+    state, grad = list_guard(state, grad)
+    beta = scalar_guard(beta, state[0])
+    _compilable_exp_avg_(state, grad, beta)
+    return grad
+
+
 @decorator_knowngood
 def _compilable_agc_(parameters: List[Tensor], gradients: List[Tensor], clip_val: float, minimum: float, eps: float):
     p_norm = torch._foreach_norm(parameters)
@@ -321,13 +336,13 @@ def nesterov_momentum(state, grad, beta):
 
 @decorator_knowngood
 def inplace_orthogonal_(x, mode, out):
-    if mode == 'qr':
+    if mode == 'newtonschulz' or x.shape[0] != x.shape[1]:
+        y = zeropower_via_newtonschulz5(x, 5)
+    elif mode == 'qr':
         y = torch.linalg.qr(x).Q
     elif mode == 'svd':
         u, s, v = torch.linalg.svd(x)
         y = u @ v.T
-    elif mode == 'newtonschulz':
-        y = zeropower_via_newtonschulz5(x, 5)
     else:
         raise NotImplementedError(f"Unknown zeroth_power_mode: {mode}")
     set_(out, y)
@@ -363,7 +378,7 @@ def get_orthogonal_matrix_QR(GG, Q, exp_avg_sq):
         est_eig = torch.einsum('ij,ij->j', o, tmp)
         sort_idx = torch.argsort(est_eig, descending=True)
         indices.append(sort_idx)
-        inplace_orthogonal_(tmp[:, sort_idx], q)
+        inplace_orthogonal_(tmp[:, sort_idx], zeroth_power_mode, q)
 
     indices = tuple(slice(None) if ind is None else ind.view(*(1,) * i, -1, *(1,) * (exp_avg_sq.dim() - i - 1))  #
                     for i, ind in enumerate(indices))

{heavyball-1.1.0 → heavyball-1.1.1}/heavyball.egg-info/PKG-INFO

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

{heavyball-1.1.0 → heavyball-1.1.1}/setup.py

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

{heavyball-1.1.0 → heavyball-1.1.1}/test/test_channels_last.py

@@ -17,7 +17,7 @@ config.cache_size_limit = 128
 
 @pytest.mark.parametrize("opt", heavyball.__all__)
 @pytest.mark.parametrize("size,depth", [(128, 1)])
-def test_foreach(opt, size, depth: int, iterations: int = 32, outer_iterations: int = 1):
+def test_foreach(opt, size, depth: int, iterations: int = 1024, outer_iterations: int = 1):
     set_torch()
     opt = getattr(heavyball, opt)
 
@@ -34,7 +34,7 @@ def test_foreach(opt, size, depth: int, iterations: int = 32, outer_iterations:
             if is_channels_last:
                 model.to(memory_format=torch.channels_last)
 
-            o = get_optim(opt, model.parameters(), lr=1e-3, weight_decay=1e-4, warmup_steps=16)
+            o = get_optim(opt, model.parameters(), lr=1e-5, weight_decay=1e-4, warmup_steps=16)
 
             for _ in range(iterations):
                 loss = model(torch.randn((1024, size, 4, 4), device='cuda')).square().mean()