heavyball 1.5.1__tar.gz → 1.5.2__tar.gz

{heavyball-1.5.1 → heavyball-1.5.2}/PKG-INFO

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

{heavyball-1.5.1 → heavyball-1.5.2}/heavyball/__init__.py

@@ -163,41 +163,6 @@ class OrthoLaProp(C.BaseOpt):
                          C.orthogonalize_grad_to_param, C.scale_by_laprop)
 
 
-class ForeachAdamW(C.BaseOpt):
-    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0,
-                 foreach: bool = True, storage_dtype: str = 'float32', mars: bool = False, caution: bool = False,
-                 mars_gamma: float = 0.0025, gradient_clipping: C.str_or_fn = C.use_default,
-                 update_clipping: C.str_or_fn = C.use_default, palm: bool = C.use_default, beta2_scale: float = 0.8):
-        defaults = locals()
-        defaults.pop("self")
-        params = defaults.pop("params")
-        super().__init__(params, defaults, foreach, gradient_clipping, update_clipping, palm, C.update_by_adam)
-
-
-class OrthoAdamW(C.BaseOpt):
-    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0,
-                 foreach: bool = True, storage_dtype: str = 'float32', mars: bool = False, caution: bool = False,
-                 mars_gamma: float = 0.0025, gradient_clipping: C.str_or_fn = C.use_default,
-                 update_clipping: C.str_or_fn = C.use_default, palm: bool = C.use_default, beta2_scale: float = 0.8):
-        defaults = locals()
-        defaults.pop("self")
-        params = defaults.pop("params")
-        super().__init__(params, defaults, foreach, gradient_clipping, update_clipping, palm,
-                         C.orthogonalize_grad_to_param, C.scale_by_adam)
-
-
-class AdamWOrtho(C.BaseOpt):
-    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0,
-                 foreach: bool = True, storage_dtype: str = 'float32', mars: bool = False, caution: bool = False,
-                 mars_gamma: float = 0.0025, gradient_clipping: C.str_or_fn = C.use_default,
-                 update_clipping: C.str_or_fn = C.use_default, palm: bool = C.use_default, beta2_scale: float = 0.8):
-        defaults = locals()
-        defaults.pop("self")
-        params = defaults.pop("params")
-        super().__init__(params, defaults, foreach, gradient_clipping, update_clipping, palm, C.scale_by_adam,
-                         C.orthogonalize_grad_to_param)
-
-
 class ForeachPSGDKron(C.BaseOpt):
     """
     Originally from Evan Walters and Omead Pooladzandi, 2024
@@ -216,7 +181,7 @@ class ForeachPSGDKron(C.BaseOpt):
                  stochastic_schedule: bool = True, storage_dtype: str = 'float32', mars: bool = False,
                  caution: bool = False, mars_gamma: float = 0.0025, delayed: Optional[bool] = C.use_default,
                  cached: Optional[bool] = C.use_default, exp_avg_input: Optional[bool] = C.use_default,
-                 gradient_clipping: C.str_or_fn = C.use_default, update_clipping: C.str_or_fn = C.use_default,  #
+                 gradient_clipping: C.str_or_fn = C.use_default, update_clipping: C.str_or_fn = C.use_default, #
                  # expert parameters
                  precond_init_scale=1.0, precond_lr=0.1):
         defaults = locals()

{heavyball-1.5.1 → heavyball-1.5.2}/heavyball/chainable.py

@@ -364,10 +364,12 @@ def _update_psgd_cache(cached, Q_cache, q):
     return Q_cache
 
 
-def _cached_psgd_precond_grad(group, cache_expr, exprs, update, Q_mat, Q_cache):
+def _cached_psgd_precond_grad(group, cache_expr, exprs, update, Q_mat, Q_cache, grad):
     if group.get('is_cached', False):
-        return utils.precond_grad_cached_(cache_expr, update, *Q_cache)
-    return utils.psgd_precond_grad(exprs[-1], update, *Q_mat)
+        out = utils.precond_grad_cached_(cache_expr, update, *Q_cache, caution=group['caution'], grad=grad)
+    out = utils.psgd_precond_grad(exprs[-1], update, *Q_mat, caution=group['caution'], grad=grad)
+    group['caution'] = False  # we already cautioned here - shouldn't do it again
+    return out
 
 
 def _fused_cached_psgd_precond_grad(group, grad, param, cache_expr, exprs, update, Q_mat, Q_cache):
@@ -387,7 +389,7 @@ def scale_by_psgd(group, update, grad, param, Q, exprs, Q_cache, cache_expr: str
     Q_mat = utils.line_to_triu(Q) if group['store_triu_as_line'] else Q
     Q_mat = _update_psgd_precond(cached, Q_cache, group, param,
                                  update if group['momentum_into_precond_update'] else grad, Q_mat, Q, exprs, prob)
-    return _cached_psgd_precond_grad(group, cache_expr, exprs, update, Q_mat, Q_cache)
+    return _cached_psgd_precond_grad(group, cache_expr, exprs, update, Q_mat, Q_cache, grad)
 
 
 @general_guard("Q", "exprs", ("Q_cache", None), ("cache_expr", None), init_fn=_init_psgd, skip_first=False)
@@ -395,7 +397,7 @@ def scale_by_psgd(group, update, grad, param, Q, exprs, Q_cache, cache_expr: str
 def scale_by_delayed_psgd(group, update, grad, param, Q, exprs, Q_cache, cache_expr: str, cached: bool = False,
                           prob: Optional[callable] = None):
     Q_mat = utils.line_to_triu(Q) if group['store_triu_as_line'] else Q
-    precond = _cached_psgd_precond_grad(group, cache_expr, exprs, update, Q_mat, Q_cache)
+    precond = _cached_psgd_precond_grad(group, cache_expr, exprs, update, Q_mat, Q_cache, grad)
     _ = _update_psgd_precond(cached, Q_cache, group, param, update if group['momentum_into_precond_update'] else grad,
                              Q_mat, Q, exprs, prob)
     return precond
@@ -467,6 +469,8 @@ class ChainOpt(utils.StatefulOptimizer):
                             f'only supported with foreach=True (currently foreach={group["foreach"]}).')
             group['base_lr'] = group['lr']
 
+        caution = group['caution']
+
         vals = list(self.split_p_and_g_in_group(group, should_promote=self.promote, beta1=utils.get_beta1(group)))
 
         if not vals:
@@ -492,6 +496,7 @@ class ChainOpt(utils.StatefulOptimizer):
         else:
             chain(self.state_, group, g, p, *self.fns)
 
+        group['caution'] = caution
         group['lr'] = group['prev_lr']
         group['step'] = None
 

{heavyball-1.5.1 → heavyball-1.5.2}/heavyball/utils.py

@@ -1300,7 +1300,10 @@ def psgd_should_update(group, prob: Union[float, callable], rng: Optional[random
 
 
 @decorator_knowngood
-def precond_grad_cached_(expr: str, ea: Tensor, *cached_q: Tensor, cast: bool = True):
+def precond_grad_cached_(expr: str, ea: Tensor, *cached_q: Tensor, caution: bool = False, grad: Optional[Tensor] = None,
+                         cast: bool = True):
+    if caution:
+        ea = _compilable_cautioning(grad, ea)
     md = min_dtype(list(cached_q) + [ea])
     args = [q.to(md) for q in cached_q]
     args = args + [ea.to(md)]
@@ -1312,8 +1315,8 @@ def precond_grad_cached_(expr: str, ea: Tensor, *cached_q: Tensor, cast: bool =
 
 @decorator_knowngood
 def _compilable_fused_precond_grad_cached_(expr: str, ea: Tensor, param, lr, grad, decay, caution, *cached_q: Tensor):
-    precond = precond_grad_cached_(expr, ea, *cached_q, cast=False)
-    update_param_(param, precond, lr, decay, caution=caution, grad=grad)
+    precond = precond_grad_cached_(expr, ea, *cached_q, caution=caution, grad=grad, cast=False)
+    update_param_(param, precond, lr, decay, caution=False)
 
 
 def fused_precond_grad_cached_(expr: str, ea: Tensor, param, lr, grad, decay, caution, *cached_q: Tensor):
@@ -1322,7 +1325,9 @@ def fused_precond_grad_cached_(expr: str, ea: Tensor, param, lr, grad, decay, ca
 
 
 @decorator_knowngood
-def psgd_precond_grad(expr: str, ea: Tensor, *preconds: Tensor):
+def psgd_precond_grad(expr: str, ea: Tensor, *preconds: Tensor, caution: bool = False, grad: Optional[Tensor] = None):
+    if caution:
+        ea = _compilable_cautioning(grad, ea)
     md = min_dtype(list(preconds) + [ea])
     args = [q.to(md) for q in preconds]
     args = args + args + [ea.to(md)]
@@ -1332,8 +1337,8 @@ def psgd_precond_grad(expr: str, ea: Tensor, *preconds: Tensor):
 
 @decorator_knowngood
 def _compilable_fused_psgd_precond_grad(expr: str, ea: Tensor, param, lr, grad, decay, caution, *preconds: Tensor):
-    precond = psgd_precond_grad(expr, ea, *preconds)
-    update_param_(param, precond, lr, decay, caution=caution, grad=grad)
+    precond = psgd_precond_grad(expr, ea, *preconds, caution=caution, grad=grad)
+    update_param_(param, precond, lr, decay, caution=False, grad=grad)
 
 
 def fused_psgd_precond_grad(expr: str, ea: Tensor, param, lr, grad, decay, caution, *preconds: Tensor):

{heavyball-1.5.1 → heavyball-1.5.2}/heavyball.egg-info/PKG-INFO

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

{heavyball-1.5.1 → heavyball-1.5.2}/setup.py

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