heavyball 0.18.0__tar.gz → 0.18.1__tar.gz

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

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

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

@@ -7,10 +7,9 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 from typing import Optional
 
 import torch
-from heavyball.utils import einsum_base
 
-from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust_region_clip_, PSGDBase, \
-    precond_update_prob_schedule, 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 ForeachCachedDelayedPSGDKron(PSGDBase):
@@ -42,7 +41,9 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,  #
+                 # expert parameters
+                 precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -55,14 +56,11 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
 
         defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
-                        momentum_into_precond_update=momentum_into_precond_update, precond_lr=0.1,
-                        # precond lr hardcoded to 0.1
-                        precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
-                        step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
-                        store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
+                        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)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
-
     def _step(self, group):
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']

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

@@ -7,10 +7,9 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 from typing import Optional
 
 import torch
-from heavyball.utils import einsum_base
 
-from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust_region_clip_, PSGDBase, \
-    precond_update_prob_schedule, 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):
@@ -40,7 +39,9 @@ class ForeachCachedPSGDKron(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,  #
+                 # expert parameters
+                 precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -53,15 +54,11 @@ class ForeachCachedPSGDKron(PSGDBase):
 
         defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
-                        momentum_into_precond_update=momentum_into_precond_update, precond_lr=0.1,
-                        # precond lr hardcoded to 0.1
-                        precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
-                        step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
-                        store_triu_as_line=store_triu_as_line,
-                        q_dtype=q_dtype)
+                        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)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
-
     def _step(self, group):
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']
@@ -119,7 +116,8 @@ class ForeachCachedPSGDKron(PSGDBase):
             if self.should_update(group):
                 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], store_triu_as_line)
+                self.do_update(group, [p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig],
+                               store_triu_as_line)
                 for c_, q_ in zip(cached_q, q):
                     if q_.ndim == 2:
                         torch.matmul(q_.T.conj(), q_, out=c_)

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

@@ -39,7 +39,9 @@ class ForeachDelayedPSGD(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: callable = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,  #
+                 # expert parameters
+                 precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -52,9 +54,8 @@ class ForeachDelayedPSGD(PSGDBase):
 
         defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
-                        momentum_into_precond_update=momentum_into_precond_update, precond_lr=0.1,
-                        # precond lr hardcoded to 0.1
-                        precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
+                        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)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)

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

@@ -5,10 +5,10 @@ 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 .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, PSGDBase, precond_update_prob_schedule, \
-    exp_avg_sq_, beta_debias, split_p_and_g_in_group, promote
+from heavyball.utils import triu_to_line, line_to_triu, identity
+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
 
 
 class ForeachPaLMPAdam(PSGDBase):
@@ -39,7 +39,9 @@ class ForeachPaLMPAdam(PSGDBase):
                  momentum_into_precond_update=True, warmup_steps: int = 1, betas=(None, None), beta: float = 0.9,
                  beta2_scale: float = 0.8, merge_dims: bool = False, split: bool = False, clip_fn: callable = None,
                  store_triu_as_line: bool = True, foreach: bool = True, q_dtype='float32',
-                 stochastic_schedule: bool = True):
+                 stochastic_schedule: bool = True,  #
+                 # expert parameters
+                 precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= weight_decay:
@@ -52,11 +54,10 @@ class ForeachPaLMPAdam(PSGDBase):
 
         defaults = dict(lr=lr, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
-                        momentum_into_precond_update=momentum_into_precond_update, precond_lr=0.1,
-                        # precond lr hardcoded to 0.1
-                        precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
-                        step=0, warmup_steps=warmup_steps, beta=beta, beta2_scale=beta2_scale, merge_dims=merge_dims,
-                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
+                        momentum_into_precond_update=momentum_into_precond_update, precond_lr=precond_lr,
+                        precond_init_scale=precond_init_scale, step=0, warmup_steps=warmup_steps, beta=beta,
+                        beta2_scale=beta2_scale, merge_dims=merge_dims, split=split,
+                        store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):

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

@@ -9,7 +9,7 @@ from typing import Optional
 import torch
 
 from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust_region_clip_, PSGDBase, \
-    precond_update_prob_schedule, split_p_and_g_in_group, line_to_triu, triu_to_line, set_, promote
+    split_p_and_g_in_group, line_to_triu, triu_to_line, set_, promote
 
 
 class ForeachPSGDKron(PSGDBase):
@@ -39,7 +39,9 @@ class ForeachPSGDKron(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,  #
+                 # expert parameters
+                 precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -52,14 +54,11 @@ class ForeachPSGDKron(PSGDBase):
 
         defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
-                        momentum_into_precond_update=momentum_into_precond_update, precond_lr=0.1,
-                        # precond lr hardcoded to 0.1
-                        precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
-                        step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
-                        store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
+                        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)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
-
     def _step(self, group):
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']
@@ -104,7 +103,7 @@ class ForeachPSGDKron(PSGDBase):
 
             if self.should_update(group):
                 q32 = [promote(q_) for q_ in q]
-                self.do_update(group,[p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
+                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))
 
         grad_list = self.clip_fn(grad_list)

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

@@ -5,10 +5,10 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 """
 
 import torch
-from heavyball.utils import copy_stochastic_list_, identity
 
-from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, PSGDBase, precond_update_prob_schedule, \
-    split_p_and_g_in_group, line_to_triu, triu_to_line, promote
+from heavyball.utils import identity
+from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, PSGDBase, split_p_and_g_in_group, \
+    line_to_triu, triu_to_line, promote
 
 
 class ForeachPurePSGD(PSGDBase):
@@ -37,8 +37,10 @@ class ForeachPurePSGD(PSGDBase):
     def __init__(self, params, lr=0.001, weight_decay=0.0, preconditioner_update_probability=None,
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
-                 split: bool = False, clip_fn: callable = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
+                 split: bool = False, clip_fn: callable = None, store_triu_as_line: bool = True, foreach: bool = True,
+                 q_dtype='float32', stochastic_schedule: bool = True,  #
+                 # expert parameters
+                 precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= weight_decay:
@@ -49,11 +51,9 @@ class ForeachPurePSGD(PSGDBase):
 
         defaults = dict(lr=lr, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
-                        momentum_into_precond_update=momentum_into_precond_update, precond_lr=0.1,
-                        # precond lr hardcoded to 0.1
-                        precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
-                        step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
-                        store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
+                        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)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
@@ -95,7 +95,7 @@ class ForeachPurePSGD(PSGDBase):
 
             if self.should_update(group):
                 q32 = [promote(q_) for q_ in q]
-                self.do_update(group,[p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
+                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)
 
         grad_list = self.clip_fn(grad_list)

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

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

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

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

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

@@ -17,7 +17,7 @@ def get_memory():
 
 @pytest.mark.parametrize("opt", heavyball.__all__)
 @pytest.mark.parametrize("size,depth", [(128, 2)])
-def test_foreach(opt, size, depth: int, iterations: int = 1024, outer_iterations: int = 3):
+def test_foreach(opt, size, depth: int, iterations: int = 8192, outer_iterations: int = 3):
     set_torch()
 
     opt = getattr(heavyball, opt)