heavyball 0.15.1__py3-none-any.whl → 0.16.0__py3-none-any.whl

heavyball/cached_psgd_kron.py

@@ -39,7 +39,8 @@ class ForeachCachedPSGDKron(PSGDBase):
     def __init__(self, params, lr=0.001, beta=0.9, 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: Optional[callable] = None, store_triu_as_line: bool = True):
+                 split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
+                 foreach: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -61,7 +62,7 @@ class ForeachCachedPSGDKron(PSGDBase):
                         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)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
         self._prob_step = 0
 

heavyball/delayed_psgd.py

@@ -38,7 +38,8 @@ class ForeachDelayedPSGD(PSGDBase):
     def __init__(self, params, lr=0.001, beta=0.9, 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):
+                 split: bool = False, clip_fn: callable = None, store_triu_as_line: bool = True,
+                 foreach: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -60,7 +61,7 @@ class ForeachDelayedPSGD(PSGDBase):
                         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)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
         self._prob_step = 0
 
@@ -113,7 +114,8 @@ class ForeachDelayedPSGD(PSGDBase):
             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 do_update:
-                self.do_update([p], [ea if momentum_into_precond_update else g], [q], precond_lr, [q_orig] if store_triu_as_line else None)
+                self.do_update([p], [ea if momentum_into_precond_update else g], [q], precond_lr,
+                               [q_orig] if store_triu_as_line else None)
                 self.balance([g], [q])
             set_(g, new)
 

heavyball/foreach_adamw.py

@@ -5,10 +5,11 @@ from .utils import warmup, exp_avg_sq_, beta_debias, update_param_, StatefulOpti
 
 
 class ForeachAdamW(StatefulOptimizer):
-    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0):
+    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0,
+                 foreach: bool = True):
         defaults = dict(lr=lr, betas=betas, eps=eps, k=0, warmup_steps=warmup_steps, train_mode=True, weight_sum=0.0,
                         lr_max=-1.0, weight_decay=weight_decay)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
     def _step(self, group):
         eps = group['eps']

heavyball/foreach_adopt.py

@@ -6,10 +6,11 @@ from .utils import warmup, beta_debias, update_param_, StatefulOptimizer
 
 class ForeachADOPT(StatefulOptimizer):
 
-    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0):
+    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0,
+                 foreach: bool = True):
         defaults = dict(lr=lr, betas=betas, eps=eps, k=0, warmup_steps=warmup_steps, train_mode=True, weight_sum=0.0,
                         lr_max=-1.0, weight_decay=weight_decay)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
     def _step(self, group):
         eps = group['eps']

heavyball/foreach_laprop.py

@@ -6,10 +6,11 @@ from .utils import warmup, exp_avg_sq_, beta_debias, update_param_, StatefulOpti
 
 class ForeachLaProp(StatefulOptimizer):
 
-    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=1):
+    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=1,
+                 foreach: bool = True):
         defaults = dict(lr=lr, betas=betas, eps=eps, k=0, warmup_steps=warmup_steps, train_mode=True, weight_sum=0.0,
                         lr_max=-1.0, weight_decay=weight_decay)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
     def _step(self, group):
         eps = group['eps']

heavyball/foreach_sfadamw.py

@@ -6,12 +6,12 @@ from .utils import schedule_free_, warmup, ScheduleFree, exp_avg_sq_, beta_debia
 
 class ForeachSFAdamW(ScheduleFree):
     def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0, r=0.0,
-                 weight_lr_power=2.0, foreach=hasattr(torch, "_foreach_mul_")):
+                 weight_lr_power=2.0, foreach: bool = True):
 
         defaults = dict(lr=lr, betas=betas, eps=eps, r=r, k=0, warmup_steps=warmup_steps, train_mode=True,
                         weight_sum=0.0, lr_max=-1.0, weight_lr_power=weight_lr_power, weight_decay=weight_decay,
                         foreach=foreach)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
     def _step(self, group):
         eps = group['eps']
@@ -48,7 +48,7 @@ class ForeachSFAdamW(ScheduleFree):
             torch._foreach_add_(grad, y, alpha=decay)
 
         lr = warmup(group['lr'], k + 1, group['warmup_steps'])
-        group['weight_sum'] = schedule_free_(lr, group['weight_lr_power'], group['weight_sum'], group['betas'][0],
-                                             y, z, grad, group['r'], k + 1)
+        group['weight_sum'] = schedule_free_(lr, group['weight_lr_power'], group['weight_sum'], group['betas'][0], y, z,
+                                             grad, group['r'], k + 1)
 
         group['k'] = k + 1

heavyball/foreach_soap.py

@@ -26,12 +26,13 @@ class ForeachSOAP(StatefulOptimizer):
                  weight_decay: float = 0.01, precondition_frequency: int = 2, max_precond_dim: int = 2048,  #
                  merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
                  data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1,
-                 split: bool = False):
+                 split: bool = False,
+                 foreach: bool = True):
         defaults = {"lr": lr, "betas": betas, "shampoo_beta": shampoo_beta, "eps": eps, "weight_decay": weight_decay,
                     "precondition_frequency": precondition_frequency, "max_precond_dim": max_precond_dim,
                     "merge_dims": merge_dims, "precondition_1d": precondition_1d, "normalize_grads": normalize_grads,
                     "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'split': split}
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
         self._data_format = data_format
 
     def _step(self, group):
@@ -59,7 +60,7 @@ class ForeachSOAP(StatefulOptimizer):
             vals.append((p, g, grad_projected, exp_avg, exp_avg_sq))
 
         if not vals:
-            return 
+            return
 
         p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
         beta1, beta2 = group["betas"]

heavyball/p_adam.py

@@ -38,7 +38,8 @@ class ForeachPaLMPAdam(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  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):
+                 store_triu_as_line: bool = True,
+                 foreach: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= weight_decay:
@@ -60,7 +61,7 @@ class ForeachPaLMPAdam(PSGDBase):
                         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)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
         self._prob_step = 0
 
@@ -90,7 +91,7 @@ class ForeachPaLMPAdam(PSGDBase):
                 state['exp_avg'] = torch.zeros_like(g)
                 state['exp_avg_sq'] = torch.zeros_like(g)
                 Q, state["exprs"] = init_Q_exprs(p, precond_init_scale, max_size_triangular,
-                                                          min_ndim_triangular, memory_save_mode, dtype=g.dtype)
+                                                 min_ndim_triangular, memory_save_mode, dtype=g.dtype)
                 state['Q'] = triu_to_line(Q) if store_triu_as_line else Q
 
             vals.append((p, g, state["Q"], state['exp_avg'], state['exp_avg_sq']))

heavyball/palm_foreach_sfadamw.py

@@ -6,13 +6,14 @@ from .utils import schedule_free_, warmup, ScheduleFree, exp_avg_sq_, beta_debia
 
 class PaLMForeachSFAdamW(ScheduleFree):
     def __init__(self, params, lr=0.0025, beta=0.9, betas=(None, None), eps=1e-8, weight_decay=0, warmup_steps=0, r=0.0,
-                 weight_lr_power=2.0, beta2_scale: float = 0.8):
+                 weight_lr_power=2.0, beta2_scale: float = 0.8,
+                 foreach: bool = True):
         if betas[0] is not None:
             beta = betas[0]
         defaults = dict(lr=lr, beta=beta, eps=eps, r=r, k=0, warmup_steps=warmup_steps, train_mode=True, weight_sum=0.0,
                         lr_max=-1.0, weight_lr_power=weight_lr_power, weight_decay=weight_decay,
                         beta2_scale=beta2_scale)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
     def _step(self, group):
         eps = group['eps']

heavyball/palm_foreach_soap.py

@@ -32,7 +32,8 @@ class PaLMForeachSOAP(StatefulOptimizer):
                  max_precond_dim: int = 2048,  #
                  merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
                  data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1,
-                 beta2_scale: float = 0.8, split: bool = False):
+                 beta2_scale: float = 0.8, split: bool = False,
+                 foreach: bool = True):
         if betas[0] is not None:
             beta = betas[0]
         defaults = {"lr": lr, "beta": beta, "shampoo_beta": shampoo_beta, "eps": eps, "weight_decay": weight_decay,
@@ -40,7 +41,7 @@ class PaLMForeachSOAP(StatefulOptimizer):
                     "merge_dims": merge_dims, "precondition_1d": precondition_1d, "normalize_grads": normalize_grads,
                     "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'beta2_scale': beta2_scale,
                     'split': split}
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
         self._data_format = data_format
 
     def _step(self, group):

heavyball/precond_schedule_foreach_soap.py

@@ -27,13 +27,14 @@ class PrecondScheduleForeachSOAP(StatefulOptimizer):
                  weight_decay: float = 0.01, precondition_frequency: int = 2, max_precond_dim: int = 2048,  #
                  merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
                  data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1,
-                 precond_scheduler=(1 / 3, 9), split: bool = False):
+                 precond_scheduler=(1 / 3, 9), split: bool = False,
+                 foreach: bool = True):
         defaults = {"lr": lr, "betas": betas, "shampoo_beta": shampoo_beta, "eps": eps, "weight_decay": weight_decay,
                     "precondition_frequency": precondition_frequency, "max_precond_dim": max_precond_dim,
                     "merge_dims": merge_dims, "precondition_1d": precondition_1d, "normalize_grads": normalize_grads,
                     "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'precond_scheduler': precond_scheduler,
                     'split': split}
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
         self._data_format = data_format
         self.rng = random.Random(0x120983109)
 

heavyball/precond_schedule_palm_foreach_soap.py

@@ -32,7 +32,8 @@ class PrecondSchedulePaLMForeachSOAP(StatefulOptimizer):
                  weight_decay: float = 0.01, precondition_frequency: int = 2, max_precond_dim: int = 2048,  #
                  merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
                  data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1,
-                 precond_scheduler=(1 / 3, 9), betas=(None, None), beta2_scale: float = 0.8, split: bool = False):
+                 precond_scheduler=(1 / 3, 9), betas=(None, None), beta2_scale: float = 0.8, split: bool = False,
+                 foreach: bool = True):
         if betas[0] is not None:
             beta = betas[0]
         defaults = {"lr": lr, "beta": beta, "shampoo_beta": shampoo_beta, "eps": eps, "weight_decay": weight_decay,
@@ -40,7 +41,7 @@ class PrecondSchedulePaLMForeachSOAP(StatefulOptimizer):
                     "merge_dims": merge_dims, "precondition_1d": precondition_1d, "normalize_grads": normalize_grads,
                     "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'precond_scheduler': precond_scheduler,
                     'beta2_scale': beta2_scale, 'split': split}
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
         self._data_format = data_format
         self.rng = random.Random(0x120983109)
 

heavyball/precond_schedule_sfpsoap.py

@@ -41,7 +41,7 @@ class PrecondScheduleSFPaLMSOAP(ScheduleFree):
                  merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
                  data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1, r=0.0,
                  weight_lr_power=2.0, gradient_clip_val: float = 0.1, precond_scheduler=(1 / 3, 9),
-                 betas=(None, None), split: bool = False):
+                 betas=(None, None), split: bool = False, foreach: bool = True):
         if betas[0] is not None:
             beta = betas[0]
         defaults = {"lr": lr, "beta": beta, "beta2_scale": beta2_scale, "eps": eps, "weight_decay": weight_decay,
@@ -50,7 +50,7 @@ class PrecondScheduleSFPaLMSOAP(ScheduleFree):
                     "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'r': r,
                     'weight_lr_power': weight_lr_power, 'train_mode': True, 'step': -1, 'weight_sum': 0,
                     'gradient_clip_val': gradient_clip_val, 'precond_scheduler': precond_scheduler, 'split': split}
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
         self._data_format = data_format
         self.rng = random.Random(0x120983109)
 
@@ -59,7 +59,7 @@ class PrecondScheduleSFPaLMSOAP(ScheduleFree):
         max_precond_dim = group['max_precond_dim']
         precondition_1d = group['precondition_1d']
 
-        step = group['step'] = group.get("step", -1) + 1
+        step = group['step'] = group.get("step", 0) + 1
 
         for p in group["params"]:
             if p.grad is None:

heavyball/psgd_kron.py

@@ -38,7 +38,8 @@ class ForeachPSGDKron(PSGDBase):
     def __init__(self, params, lr=0.001, beta=0.9, 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: Optional[callable] = None, store_triu_as_line: bool = True):
+                 split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
+                 foreach: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -60,7 +61,7 @@ class ForeachPSGDKron(PSGDBase):
                         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)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
         self._prob_step = 0
 
@@ -114,7 +115,8 @@ class ForeachPSGDKron(PSGDBase):
 
             if do_update:
                 self.balance([g], [q])
-                self.do_update([p], [ea if momentum_into_precond_update else g], [q], precond_lr, [q_orig] if store_triu_as_line else None)
+                self.do_update([p], [ea if momentum_into_precond_update else g], [q], precond_lr,
+                               [q_orig] if store_triu_as_line else None)
             set_(g, psgd_precond_grad(q, self.state_(p)["exprs"], ea))
 
         grad_list = self.clip_fn(grad_list)

heavyball/pure_psgd.py

@@ -36,7 +36,8 @@ 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):
+                 split: bool = False, clip_fn: callable = None, store_triu_as_line: bool = True,
+                 foreach: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= weight_decay:
@@ -56,7 +57,7 @@ class ForeachPurePSGD(PSGDBase):
                         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)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
 
         self._prob_step = 0
 

heavyball/schedule_free_palm_foreach_soap.py

@@ -33,7 +33,8 @@ class SFPaLMForeachSOAP(ScheduleFree):
                  weight_decay: float = 0.01, precondition_frequency: int = 2, max_precond_dim: int = 2048,  #
                  merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
                  data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1, r=0.0,
-                 weight_lr_power=2.0, gradient_clip_val: float = 0.1, betas=(None, None), split: bool = False):
+                 weight_lr_power=2.0, gradient_clip_val: float = 0.1, betas=(None, None), split: bool = False,
+                 foreach: bool = True):
         if betas[0] is not None:
             beta = betas[0]
         defaults = {"lr": lr, "beta": beta, "beta2_scale": beta2_scale, "eps": eps, "weight_decay": weight_decay,
@@ -42,7 +43,7 @@ class SFPaLMForeachSOAP(ScheduleFree):
                     "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'r': r,
                     'weight_lr_power': weight_lr_power, 'train_mode': True, 'step': -1,
                     'gradient_clip_val': gradient_clip_val, 'weight_sum': 0, 'split': split}
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, foreach)
         self._data_format = data_format
         self.rng = random.Random(0x120983109)
 
@@ -51,7 +52,7 @@ class SFPaLMForeachSOAP(ScheduleFree):
         max_precond_dim = group['max_precond_dim']
         precondition_1d = group['precondition_1d']
 
-        step = group['step'] = group.get("step", -1) + 1
+        step = group['step'] = group.get("step", 0) + 1
 
         for p in group["params"]:
             if p.grad is None:

heavyball/utils.py

@@ -383,8 +383,25 @@ def project(grad, Q, back: bool):
 
 
 class StatefulOptimizer(torch.optim.Optimizer):
+    def __init__(self, params, defaults, foreach: bool = True):
+        super().__init__(params, {**defaults, 'foreach': foreach})
+        self.fake_groups = {}
+
+    def key(self, param: torch.Tensor):
+        return (param.data_ptr(), tuple(param.shape))
+
+    def get_groups(self, group):
+        if group['foreach']:
+            return [group]
+
+        for p in group['params']:
+            if self.key(p) not in self.fake_groups:
+                self.fake_groups[self.key(p)] = {**group, 'params': [p]}
+
+        return [self.fake_groups[self.key(p)] for p in group['params']]
+
     def state_(self, arg: torch.Tensor):
-        return self.state[(arg.data_ptr(), tuple(arg.shape))]
+        return self.state[self.key(arg)]
 
     def state_size(self) -> int:
         total_bytes = 0
@@ -409,8 +426,9 @@ class StatefulOptimizer(torch.optim.Optimizer):
             with torch.enable_grad():
                 loss = closure()
         with torch.no_grad():
-            for group in self.param_groups:
-                self._step(group)
+            for top_group in self.param_groups:
+                for group in self.get_groups(top_group):
+                    self._step(group)
         return loss
 
 
@@ -754,8 +772,8 @@ def update_triu_(q_state, materialised):
 
 
 class PSGDBase(StatefulOptimizer):
-    def __init__(self, parameters, groups):
-        super().__init__(parameters, groups)
+    def __init__(self, parameters, groups, foreach: bool = True):
+        super().__init__(parameters, groups, foreach)
         self.rng = random.Random(0x1923213)
         self._tiny = torch.finfo(torch.bfloat16).tiny
 

{heavyball-0.15.1.dist-info → heavyball-0.16.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.15.1
+Version: 0.16.0
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler
@@ -39,12 +39,14 @@ recommended experimental optimizer is `ForeachPSGDKron`.
 
 * **Stochastic Rounding**: [FP32 convergence with BF16 parameters](https://github.com/pytorch/pytorch/issues/120376)
 * **Inplace EMA**: Same math, but less memory, less compute and higher stability
-* **Foreach**: Fast multi-tensor application
+* **Foreach**: Fast multi-tensor application (turn it off to save memory via `foreach=False`)
 * **PaLM Beta2**: Fast initial
   convergence, [stable late convergence](https://x.com/_clashluke/status/1820810798693818761)
 * **ScheduleFree**: No learning rate schedule, but better convergence
 * [**Preconditioner Schedule**](https://github.com/lixilinx/psgd_torch/): Improved loss-per-step in early convergence,
   better step-per-second in late convergence (explained below)
+* **Memory-efficient storage** PSGD supports `store_triu_as_line` (default: `True`) to trade off memory usage for memory
+  bandwidth; turn it off for lower overheads (for more, see [PSGD Efficiency](docs/psgd_efficiency.md))
 
 ## Getting started
 

heavyball-0.16.0.dist-info/RECORD

@@ -0,0 +1,23 @@
+heavyball/__init__.py,sha256=KbT0GMU0DKqZxq9laCrD7XgiqS9yxC1W52zhte5kjKs,2054
+heavyball/cached_psgd_kron.py,sha256=vJuy639G-_ZLSRX3goSFMXALv-ucYjrxaEtpj0IHo-M,6802
+heavyball/delayed_psgd.py,sha256=sbwgAed5gmQpHNTPvuE7Si-gB-s0NVvN4d-4rNUJj4c,5893
+heavyball/foreach_adamw.py,sha256=h_ar0ZZRM0q_wxuEkxEOEYe-2p-mB4OMgAHivrUnPl8,1777
+heavyball/foreach_adopt.py,sha256=ogOw2JjwEQNj7AKlweAphQFdMJ_GcMDm-RyDvEzugoc,1911
+heavyball/foreach_laprop.py,sha256=yGVmGqWiSw8Y2Xj70ndkR8ZMygakTB4_iRwV02Svkqg,1816
+heavyball/foreach_sfadamw.py,sha256=15-n6-lx4PAHYsKYmXbugxsR5MnqaPYy2vUudPRiitg,2087
+heavyball/foreach_soap.py,sha256=h6ptMch7oaynvu3eIJtWnVXypDA_5JDVm3Zb3PNEma0,4634
+heavyball/p_adam.py,sha256=aCu4Qn0eHJETHuCGrfNKp2aygKk2ZoNQyxut3Vcqmoc,6112
+heavyball/palm_foreach_sfadamw.py,sha256=yvZbPyjDW8qd3r4qDXb6uTr5RozQ7JSDj4aYYRnKGLA,2248
+heavyball/palm_foreach_soap.py,sha256=g4hbiGRcti-J-a0SwAkP4ii5pU-aalsZH5bssyhroLk,5938
+heavyball/precond_schedule_foreach_soap.py,sha256=WLg5SzpJnKPZUvFyIvdwSZa1Umt5cpr3Kow_42orM-E,4863
+heavyball/precond_schedule_palm_foreach_soap.py,sha256=ammQrvRZFF-wc-wEiPEoFhS_7b8pdV61QfcLoQfimSo,6211
+heavyball/precond_schedule_sfpsoap.py,sha256=vq7jd302refKPa_9X2lkOTOtCCcTBVByPdojklrY8pA,6770
+heavyball/psgd_kron.py,sha256=iWTAViuzxTodtQGZnkLsEXrLG8tNU-BQB3KkTYAVcX4,5874
+heavyball/pure_psgd.py,sha256=EuCPNM8TX13cOop-mvvBFh6Uo1UjD1vsE053hvil92Q,5136
+heavyball/schedule_free_palm_foreach_soap.py,sha256=zkcikH5wWbzq4kOrmBjilvY3iWzuUddcv2HNEPKr3MI,6366
+heavyball/utils.py,sha256=z6taEvpgszKTrscqgowKYqb0xIVpBDVDBNGgvTE4Pb8,28484
+heavyball-0.16.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.16.0.dist-info/METADATA,sha256=yjpldOTN2rXN2-KG7R9ytuyBfmSCDpznZeRuziANChE,11941
+heavyball-0.16.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.16.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.16.0.dist-info/RECORD,,

heavyball-0.15.1.dist-info/RECORD

@@ -1,23 +0,0 @@
-heavyball/__init__.py,sha256=KbT0GMU0DKqZxq9laCrD7XgiqS9yxC1W52zhte5kjKs,2054
-heavyball/cached_psgd_kron.py,sha256=mXDtxq2WJST_aUJhrLr_xCCXSFaDvD5gCTSEveBUtac,6754
-heavyball/delayed_psgd.py,sha256=dN3NW1jmjxmUkgqxPwUVrqLY8nnBOFp4TVtJ_BhPDR4,5814
-heavyball/foreach_adamw.py,sha256=NSzoIgNm7eavzbJgkAF0k7TUEnWAgOpt9-4juIFoaSA,1729
-heavyball/foreach_adopt.py,sha256=WA07m5jocLfb1GPU8s6mJ2PteS-03ronkKm-VJrAm5I,1863
-heavyball/foreach_laprop.py,sha256=mE2NDGX9XgvRhsewcWnk_-FulZPqGA65ejYF_9-A1Xk,1768
-heavyball/foreach_sfadamw.py,sha256=ussHfPd99u3RTfMrCuu5oIbwNFLXK19wO1Fbz3JShlc,2097
-heavyball/foreach_soap.py,sha256=WWvssYKg607uoEJHftp8ag8mtKSKSeHrT0QTgqBucVg,4587
-heavyball/p_adam.py,sha256=ms7BoMHu3jKGsuztUeECrsXufGAwBpqGsxgZ5LBXLQg,6073
-heavyball/palm_foreach_sfadamw.py,sha256=wjUb_fNZNUmzWXyKvwB0unP9lvNMmaYSQo5YoeS5cj0,2200
-heavyball/palm_foreach_soap.py,sha256=2Sb4hUHQeexJcCgjHeQM_ENkZ6lG1DVxW72ryrvR6iY,5890
-heavyball/precond_schedule_foreach_soap.py,sha256=bHsDyh-UvHpHjumjqqy0PePoR1ZMsJV6o5wWvpLAA04,4815
-heavyball/precond_schedule_palm_foreach_soap.py,sha256=myLTJNQKLtZ3Xi3MVTB-RYtx_XeMRJw5CIMJW75ndUY,6163
-heavyball/precond_schedule_sfpsoap.py,sha256=xeNWetBzBEYqfOSzl98aAVJsHk43QkrUUhHH_YD_mS4,6740
-heavyball/psgd_kron.py,sha256=rMG5UPEgyfQs_n1MHSEicekVDpbbIzinlL8akEyY918,5795
-heavyball/pure_psgd.py,sha256=LLVJhUAb04hgAmT3BTz_faswwQEQUkLhm_VwGQmbBUo,5088
-heavyball/schedule_free_palm_foreach_soap.py,sha256=w0P7lMmoijTpL9V7NwOHcNBFJQ7S1TS9aCiwPhY2yVw,6319
-heavyball/utils.py,sha256=PWmwjZPL4oxMjK79a5R1e7JHykphNi5GdpYqO_xmmFU,27829
-heavyball-0.15.1.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.15.1.dist-info/METADATA,sha256=0wImMJNYM-Zg0akh9hRf7X8ofVW6zlmpyDGgAkK5GFA,11667
-heavyball-0.15.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.15.1.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.15.1.dist-info/RECORD,,