heavyball 0.18.1__tar.gz → 0.18.3__tar.gz

{heavyball-0.18.1 → heavyball-0.18.3}/PKG-INFO

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

{heavyball-0.18.1 → heavyball-0.18.3}/heavyball/cached_delayed_psgd_kron.py

@@ -62,6 +62,7 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
+        should_update = self.should_update(group)
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
@@ -113,7 +114,9 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
             q_orig = Q_list.pop(0)
             ea = exp_avg_list.pop(0)
 
-            if self.should_update(group):
+            new = torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea)
+
+            if should_update:
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig
                 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],
@@ -124,7 +127,7 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                     else:
                         torch.mul(q_.conj(), q_, out=c_)
 
-            set_(g, torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea))
+            set_(g, new)
         grad_list = self.clip_fn(grad_list)
 
         lr = -warmup(lr, group['step'], group['warmup_steps'])

{heavyball-0.18.1 → heavyball-0.18.3}/heavyball/cached_psgd_kron.py

@@ -8,8 +8,8 @@ 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, set_, einsum_base, promote
+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, set_, einsum_base, promote
 
 
 class ForeachCachedPSGDKron(PSGDBase):
@@ -71,6 +71,7 @@ class ForeachCachedPSGDKron(PSGDBase):
         beta = group['beta']
         store_triu_as_line = group['store_triu_as_line']
         q_dtype = getattr(torch, group['q_dtype'])
+        should_update = self.should_update(group)
 
         vals = []
 
@@ -111,9 +112,7 @@ class ForeachCachedPSGDKron(PSGDBase):
             q_orig = Q_list.pop(0)
             ea = exp_avg_list.pop(0)
 
-            new = torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea)
-
-            if self.should_update(group):
+            if should_update:
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig
                 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],
@@ -124,7 +123,7 @@ class ForeachCachedPSGDKron(PSGDBase):
                     else:
                         torch.mul(q_.conj(), q_, out=c_)
 
-            set_(g, new)
+            set_(g, torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea))
 
         grad_list = self.clip_fn(grad_list)
 

{heavyball-0.18.1 → heavyball-0.18.3}/heavyball/delayed_psgd.py

@@ -62,6 +62,7 @@ class ForeachDelayedPSGD(PSGDBase):
 
 
     def _step(self, group):
+        should_update = self.should_update(group)
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
@@ -103,7 +104,7 @@ class ForeachDelayedPSGD(PSGDBase):
             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)
-            if self.should_update(group):
+            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)
             set_(g, new)

{heavyball-0.18.1 → heavyball-0.18.3}/heavyball/p_adam.py

@@ -61,6 +61,7 @@ class ForeachPaLMPAdam(PSGDBase):
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
+        should_update = self.should_update(group)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
         min_ndim_triangular = group['min_ndim_triangular']
@@ -94,7 +95,7 @@ class ForeachPaLMPAdam(PSGDBase):
         group["step"] += 1
 
         Q_triu = [line_to_triu(q) if store_triu_as_line else q for q in Q_list]
-        if self.should_update(group):
+        if should_update:
             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)

{heavyball-0.18.1 → heavyball-0.18.3}/heavyball/psgd_kron.py

@@ -60,6 +60,7 @@ class ForeachPSGDKron(PSGDBase):
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
+        should_update = self.should_update(group)
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
@@ -101,7 +102,7 @@ class ForeachPSGDKron(PSGDBase):
             ea = exp_avg_list.pop(0)
             q = line_to_triu(q_orig) if store_triu_as_line else q_orig
 
-            if self.should_update(group):
+            if should_update:
                 q32 = [promote(q_) for q_ in q]
                 self.do_update(group, [p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
             set_(g, psgd_precond_grad(q, self.state_(p)["exprs"], ea))

{heavyball-0.18.1 → heavyball-0.18.3}/heavyball/pure_psgd.py

@@ -57,7 +57,7 @@ class ForeachPurePSGD(PSGDBase):
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
-        # update preconditioners all together
+        should_update = self.should_update(group)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
         min_ndim_triangular = group['min_ndim_triangular']
@@ -93,7 +93,7 @@ class ForeachPurePSGD(PSGDBase):
             q_orig = Q_list.pop(0)
             q = line_to_triu(q_orig) if store_triu_as_line else q_orig
 
-            if self.should_update(group):
+            if group:
                 q32 = [promote(q_) for q_ in q]
                 self.do_update(group, [p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
             psgd_precond_grad(q, self.state_(p)["exprs"], g, inplace=True)

{heavyball-0.18.1 → heavyball-0.18.3}/heavyball/utils.py

@@ -335,9 +335,7 @@ def promote(x):
 def min_dtype(xs: List[torch.Tensor]):
     dtypes = [x.dtype for x in xs]
     for d in (torch.float32, torch.bfloat16, torch.float16):
-        if all(d == x for x in dtypes):
-            return d
-        if all(d in (x, torch.float32, torch.float64) for x in dtypes):
+        if all(x in (d, torch.float32, torch.float64) for x in dtypes):
             return d
     return torch.float32
 
@@ -817,14 +815,15 @@ class PSGDBase(StatefulOptimizer):
 
     def do_update(self, group, p_list, grad_list, q_list, precond_lr, original_q: Optional[List] = None,
                   store_triu_as_line=False):
-        if self.should_update(group, self.balance_probability, 'balance_prob'):
-            for g, q in zip(grad_list, q_list):
-                if g.dim() > 1:
-                    psgd_balance_Q(q)
-
         for i, (p, grad, Q) in enumerate(zip(p_list, grad_list, q_list)):
             psgd_update_precond(Q, self.state_(p)["exprs"], torch.randn_like(grad), grad, precond_lr, self._tiny)
-            if original_q:
+
+        for g, q in zip(grad_list, q_list):
+            if g.dim() > 1:
+                psgd_balance_Q(q)
+
+        if original_q:
+            for q in q_list:
                 if store_triu_as_line:
                     update_triu_(original_q[i], Q)
                 else:

{heavyball-0.18.1 → heavyball-0.18.3}/heavyball.egg-info/PKG-INFO

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

{heavyball-0.18.1 → heavyball-0.18.3}/setup.py

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

{heavyball-0.18.1 → heavyball-0.18.3}/test/test_stochastic_updates.py

@@ -16,7 +16,7 @@ def get_memory():
 
 
 @pytest.mark.parametrize("opt", heavyball.__all__)
-@pytest.mark.parametrize("size,depth", [(128, 2)])
+@pytest.mark.parametrize("size,depth", [(128, 1)])
 def test_foreach(opt, size, depth: int, iterations: int = 8192, outer_iterations: int = 3):
     set_torch()
 
@@ -28,12 +28,13 @@ def test_foreach(opt, size, depth: int, iterations: int = 8192, outer_iterations
     losses = []
 
     for stochastic in [False, True]:
+        print('stochastic', stochastic)
         torch.manual_seed(0x2131290)
         peaks.append([])
         losses.append([])
 
         for i in range(outer_iterations):
-            model = nn.Sequential(*[nn.Linear(size, size) for _ in range(depth)]).cuda()
+            model = nn.Sequential(*[nn.Linear(size, size, bias=False) for _ in range(depth)]).cuda()
             o = get_optim(opt, model.parameters(), lr=1e-3, stochastic_schedule=stochastic)
 
             for _ in range(iterations):