heavyball 0.25.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

{heavyball-0.25.0.dist-info → heavyball-1.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.25.0
+Version: 1.0.0
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler
@@ -32,11 +32,12 @@ A simple package of efficient optimizers
 The goal is not to thrive for completeness, full maintenance or abstraction, but instead to provide a simple
 largely static alternative to `torch.optim` with more and better optimizers.
 
-Currently (2024-11-26, 0.22.1), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
+Currently (2024-12-07, 1.0.0), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
 recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psgd_efficiency.md)).
 
 ## Features
 
+* **Optax-like API**: `C = heavyball.chainable; grokfast = C.ChainOpt(p, lr, C.exp_avg, C.scale_by_adam)`
 * **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 (turn it off to save memory via `foreach=False`)
@@ -130,6 +131,6 @@ To access `heavyball.utils`, you need to explicitly `import heavyball.utils`.\
 It has several handy functions:
 
 * `set_torch()` sets pytorch optimization settings (TF32, opt_einsum, benchmark, ...)
-* `compile_mode`, a string passed as-is to `torch.compile(mode=compile_mode)` in all compiled heavyball calls
+* `compile_mode`, a string passed as-is to `torch.compile(mode=compile_mode)` in all compiled heavyball calls; `compile_mode=None` disables torch_compile
 * `zeroth_power_mode`, a string determining whether to use QR, newtonschulz{iterations}, or svd or eigh to approximate
   the eigenvectors. Eigh has the highest precision and cost

heavyball-1.0.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+heavyball/__init__.py,sha256=1QPYBIH8amnk3-_rKe6L9FJ0rkV5wVNRr7Yw9BXjIYI,11636
+heavyball/chainable.py,sha256=cp-tpetPr4CNN9xJ85JSo89JYC5BWUygoE6dnET6tmc,18141
+heavyball/utils.py,sha256=qUoB9EIxl7GUyLkV5a5JAKOD6TvPc1FNsqyUbJ-HY6o,46343
+heavyball-1.0.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-1.0.0.dist-info/METADATA,sha256=9C2btIxngp26TRCJFU6B8ftkWQt1rfZZC10rkAhaORw,12074
+heavyball-1.0.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-1.0.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-1.0.0.dist-info/RECORD,,

heavyball/cached_delayed_psgd_kron.py

@@ -1,135 +0,0 @@
-"""
-Originally from Evan Walters and Omead Pooladzandi, 2024
-Modified under Creative Commons Attribution 4.0 International
-Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5ee8a1a4c29e4780bbf6c521e545189eff9/kron_torch/kron.py
-"""
-
-from typing import Optional
-
-import torch
-from heavyball.utils import min_dtype, precond_grad_cached_
-
-from .utils import update_param_, warmup, init_Q_exprs, trust_region_clip_, PSGDBase,  \
-    line_to_triu, triu_to_line, einsum_base, promote, stochastic_lerp_, beta_debias, precond_grad_cached_
-
-
-class ForeachCachedDelayedPSGDKron(PSGDBase):
-    """
-    Implements PSGD with off-by-one preconditioning (akin to ADOPT and SOAP) with cached preconditioners.
-
-
-    Args:
-        params (iterable): Iterable of parameters to optimize or dicts defining
-            parameter groups.
-        lr (float): Learning rate.
-        beta (float): Momentum parameter.
-        weight_decay (float): Weight decay (L2 penalty).
-        preconditioner_update_probability (callable or float, optional): Probability of
-            updating the preconditioner. If None, defaults to a schedule that anneals
-            from 1.0 to 0.03 by 4000 steps.
-        max_size_triangular (int): Max size for dim's preconditioner to be triangular.
-        min_ndim_triangular (int): Minimum number of dimensions a layer needs
-            to have triangular preconditioners.
-        memory_save_mode: (string, optional), None, 'one_diag', or 'all_diag', None is default
-            to set all preconditioners to be triangular, 'one_diag' sets the largest
-            or last dim to be diagonal per layer, and 'all_diag' sets all preconditioners
-            to be diagonal.
-        momentum_into_precond_update: (bool), whether to send momentum into preconditioner
-            update instead of raw gradients.
-    """
-
-    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,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,
-                 storage_dtype: str = 'float32', mars: bool = False, caution: bool = False, mars_gamma: float = 0.0025,
-                 #
-                 # 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:
-            raise ValueError(f"Invalid beta parameter: {beta}")
-        if not 0.0 <= weight_decay:
-            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
-
-        if clip_fn is None:
-            clip_fn = lambda x: trust_region_clip_(x, 0.9, 1.5)
-
-        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=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,
-                        storage_dtype=storage_dtype, caution=caution, mars_gamma=mars_gamma, mars=mars)
-        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']
-        min_ndim_triangular = group['min_ndim_triangular']
-        memory_save_mode = group['memory_save_mode']
-        precond_lr = group['precond_lr']
-        weight_decay = group['weight_decay']
-        lr = group['lr']
-        beta = group['beta']
-        store_triu_as_line = group['store_triu_as_line']
-        q_dtype = getattr(torch, group['q_dtype'])
-        storage_dtype = getattr(torch, group['storage_dtype'])
-
-        vals = []
-
-        for p, g in self.split_p_and_g_in_group(group, should_promote=False, beta1=beta):
-            state = self.state_(p)
-
-            if 'Q' not in state:
-                state["exp_avg"] = torch.zeros_like(g, dtype=storage_dtype, memory_format=torch.preserve_format)
-                Q, state["exprs"] = init_Q_exprs(p, precond_init_scale, max_size_triangular, min_ndim_triangular,
-                                                 memory_save_mode, dtype=q_dtype)
-                state['Q'] = triu_to_line(Q) if store_triu_as_line else Q
-                state['Q_cache'] = [torch.empty_like(q) for q in Q]
-
-                expr = [f'{c.upper()}{c}' if q_.ndim == 2 else c for c, q_ in zip(einsum_base, Q)]
-                expr = ','.join(expr)
-                grad_expr = ''.join(c for c, _ in zip(einsum_base, g.shape))
-                out_expr = ''.join(c.upper() if c.upper() in expr else c for c in grad_expr)
-                expr = f'{expr},{grad_expr}->{out_expr}'
-
-                state['cache_expr'] = expr
-
-            vals.append((p, g, state["exp_avg"], state["Q"], state['Q_cache']))
-
-        if not vals:
-            return
-
-        p_list, grad_list, exp_avg_list, Q_list, Q_cache_list = zip(*vals)
-        del vals
-
-        group["step"] += 1
-
-        stochastic_lerp_(exp_avg_list, grad_list, 1 - beta_debias(beta, group['step']))
-
-        lr = -warmup(lr, group['step'], group['warmup_steps'])
-
-        grad_list, Q_list, Q_cache_list, exp_avg_list = list(grad_list), list(Q_list), list(Q_cache_list), list(
-            exp_avg_list)
-        for i, (p, g) in enumerate(zip(p_list, grad_list)):
-            cached_q = Q_cache_list.pop(0)
-            q_orig = Q_list.pop(0)
-            ea = exp_avg_list.pop(0)
-
-            precond_grad_cached_(cached_q, ea, self.state_(p)['cache_expr'], p, lr, weight_decay, self.clip_fn, group['caution'], g)
-
-            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],
-                               store_triu_as_line)
-                for c_, q_ in zip(cached_q, q):
-                    if q_.ndim == 2:
-                        torch.matmul(q_.T.conj(), q_, out=c_)
-                    else:
-                        torch.mul(q_.conj(), q_, out=c_)

heavyball/cached_psgd_kron.py

@@ -1,136 +0,0 @@
-"""
-Originally from Evan Walters and Omead Pooladzandi, 2024
-Modified under Creative Commons Attribution 4.0 International
-Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5ee8a1a4c29e4780bbf6c521e545189eff9/kron_torch/kron.py
-"""
-
-from typing import Optional
-
-import torch
-
-from .utils import update_param_, warmup, init_Q_exprs, trust_region_clip_, PSGDBase, \
-    line_to_triu, triu_to_line, einsum_base, promote, stochastic_lerp_, beta_debias, precond_grad_cached_
-
-
-class ForeachCachedPSGDKron(PSGDBase):
-    """Implements PSGD Kron from https://github.com/lixilinx/psgd_torch with cached preconditioners.
-
-    Args:
-        params (iterable): Iterable of parameters to optimize or dicts defining
-            parameter groups.
-        lr (float): Learning rate.
-        beta (float): Momentum parameter.
-        weight_decay (float): Weight decay (L2 penalty).
-        preconditioner_update_probability (callable or float, optional): Probability of
-            updating the preconditioner. If None, defaults to a schedule that anneals
-            from 1.0 to 0.03 by 4000 steps.
-        max_size_triangular (int): Max size for dim's preconditioner to be triangular.
-        min_ndim_triangular (int): Minimum number of dimensions a layer needs
-            to have triangular preconditioners.
-        memory_save_mode: (string, optional), None, 'one_diag', or 'all_diag', None is default
-            to set all preconditioners to be triangular, 'one_diag' sets the largest
-            or last dim to be diagonal per layer, and 'all_diag' sets all preconditioners
-            to be diagonal.
-        momentum_into_precond_update: (bool), whether to send momentum into preconditioner
-            update instead of raw gradients.
-    """
-
-    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,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,
-                 storage_dtype: str = 'float32', mars: bool = False, caution: bool = False, mars_gamma: float = 0.0025,
-                 orthogonalize_output: bool = False,
-                 #
-                 # 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:
-            raise ValueError(f"Invalid beta parameter: {beta}")
-        if not 0.0 <= weight_decay:
-            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
-
-        if clip_fn is None:
-            clip_fn = lambda x: trust_region_clip_(x, 0.9, 1.5)
-
-        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=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,
-                        storage_dtype=storage_dtype, caution=caution, mars_gamma=mars_gamma, mars=mars,
-                        orthogonalize_output=orthogonalize_output)
-        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']
-        max_size_triangular = group['max_size_triangular']
-        min_ndim_triangular = group['min_ndim_triangular']
-        memory_save_mode = group['memory_save_mode']
-        precond_lr = group['precond_lr']
-        weight_decay = group['weight_decay']
-        lr = group['lr']
-        beta = group['beta']
-        store_triu_as_line = group['store_triu_as_line']
-        q_dtype = getattr(torch, group['q_dtype'])
-        storage_dtype = getattr(torch, group['storage_dtype'])
-        orthogonalize_output = group['orthogonalize_output']
-        should_update = self.should_update(group)
-
-        vals = []
-
-        for p, g in self.split_p_and_g_in_group(group, should_promote=False, beta1=beta):
-            state = self.state_(p)
-
-            if 'Q' not in state:
-                state["exp_avg"] = torch.zeros_like(g, dtype=storage_dtype, memory_format=torch.preserve_format)
-                Q, state["exprs"] = init_Q_exprs(p, precond_init_scale, max_size_triangular, min_ndim_triangular,
-                                                 memory_save_mode, dtype=q_dtype)
-                state['Q'] = triu_to_line(Q) if store_triu_as_line else Q
-                state['Q_cache'] = [torch.empty_like(q) for q in Q]
-
-                expr = [f'{c.upper()}{c}' if q_.ndim == 2 else c for c, q_ in zip(einsum_base, Q)]
-                expr = ','.join(expr)
-                grad_expr = ''.join(c for c, _ in zip(einsum_base, g.shape))
-                out_expr = ''.join(c.upper() if c.upper() in expr else c for c in grad_expr)
-                expr = f'{expr},{grad_expr}->{out_expr}'
-
-                state['cache_expr'] = expr
-
-            vals.append((p, g, state["exp_avg"], state["Q"], state['Q_cache']))
-
-        if not vals:
-            return
-
-        p_list, grad_list, exp_avg_list, Q_list, Q_cache_list = zip(*vals)
-        del vals
-
-        group["step"] += 1
-
-        stochastic_lerp_(exp_avg_list, grad_list, 1 - beta_debias(beta, group['step']))
-
-        lr = -warmup(lr, group['step'], group['warmup_steps'])
-
-        grad_list, Q_list, Q_cache_list, exp_avg_list = list(grad_list), list(Q_list), list(Q_cache_list), list(
-            exp_avg_list)
-        for i, (p, g) in enumerate(zip(p_list, grad_list)):
-            cached_q = Q_cache_list.pop(0)
-            q_orig = Q_list.pop(0)
-            ea = exp_avg_list.pop(0)
-
-            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],
-                               store_triu_as_line)
-                for c_, q_ in zip(cached_q, q):
-                    if q_.ndim == 2:
-                        torch.matmul(q_.T.conj(), q_, out=c_)
-                    else:
-                        torch.mul(q_.conj(), q_, out=c_)
-
-            precond_grad_cached_(cached_q, ea, self.state_(p)['cache_expr'], p, lr, weight_decay, self.clip_fn,
-                                 group['caution'], g)

heavyball/delayed_psgd.py

@@ -1,122 +0,0 @@
-"""
-Originally from Evan Walters and Omead Pooladzandi, 2024
-Modified under Creative Commons Attribution 4.0 International
-Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5ee8a1a4c29e4780bbf6c521e545189eff9/kron_torch/kron.py
-"""
-
-import torch
-from heavyball.utils import stochastic_lerp_, beta_debias, stochastic_add_
-
-from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust_region_clip_, PSGDBase, \
-    triu_to_line, line_to_triu, promote,_compilable_update_, decorator_knowngood
-
-
-@decorator_knowngood
-def _compilable_psgd_precond_grad_(q, exprs, ea, p, lr, weight_decay, clip_fn, caution, grad):
-    new = psgd_precond_grad(False, exprs, ea, *q)
-    _compilable_update_([p], clip_fn([new]), weight_decay, stochastic_add_, lr, caution, [grad])
-
-
-class ForeachDelayedPSGD(PSGDBase):
-    """
-    Implements PSGD with off-by-one preconditioning (akin to ADOPT and SOAP)
-
-    Args:
-        params (iterable): Iterable of parameters to optimize or dicts defining
-            parameter groups.
-        lr (float): Learning rate.
-        beta (float): Momentum parameter.
-        weight_decay (float): Weight decay (L2 penalty).
-        preconditioner_update_probability (callable or float, optional): Probability of
-            updating the preconditioner. If None, defaults to a schedule that anneals
-            from 1.0 to 0.03 by 4000 steps.
-        max_size_triangular (int): Max size for dim's preconditioner to be triangular.
-        min_ndim_triangular (int): Minimum number of dimensions a layer needs
-            to have triangular preconditioners.
-        memory_save_mode: (string, optional), None, 'one_diag', or 'all_diag', None is default
-            to set all preconditioners to be triangular, 'one_diag' sets the largest
-            or last dim to be diagonal per layer, and 'all_diag' sets all preconditioners
-            to be diagonal.
-        momentum_into_precond_update: (bool), whether to send momentum into preconditioner
-            update instead of raw gradients.
-    """
-
-    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, foreach: bool = True,
-                 q_dtype='float32', stochastic_schedule: bool = True, storage_dtype: str = 'float32',
-                 mars: bool = False, caution: bool = False, mars_gamma: float = 0.0025,  #
-                 # 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:
-            raise ValueError(f"Invalid beta parameter: {beta}")
-        if not 0.0 <= weight_decay:
-            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
-
-        if clip_fn is None:
-            clip_fn = lambda x: trust_region_clip_(x, 0.9, 1.5)
-
-        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=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,
-                        storage_dtype=storage_dtype,
-                        caution=caution, mars_gamma=mars_gamma, mars=mars)
-        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']
-        min_ndim_triangular = group['min_ndim_triangular']
-        memory_save_mode = group['memory_save_mode']
-        precond_lr = group['precond_lr']
-        weight_decay = group['weight_decay']
-        lr = group['lr']
-        beta = group['beta']
-        store_triu_as_line = group['store_triu_as_line']
-        q_dtype = getattr(torch, group['q_dtype'])
-        storage_dtype = getattr(torch, group['storage_dtype'])
-
-        vals = []
-
-        for p, g in self.split_p_and_g_in_group(group, should_promote=False, beta1=beta):
-            state = self.state_(p)
-
-            if 'Q' not in state:
-                state["exp_avg"] = torch.zeros_like(g, dtype=storage_dtype, memory_format=torch.preserve_format)
-                Q, state["exprs"] = init_Q_exprs(p, precond_init_scale, max_size_triangular, min_ndim_triangular,
-                                                 memory_save_mode, dtype=q_dtype)
-                state["Q"] = triu_to_line(Q) if store_triu_as_line else Q
-
-            vals.append((p, g, state["exp_avg"], state["Q"]))
-
-        if not vals:
-            return
-
-        p_list, grad_list, exp_avg_list, Q_list = zip(*vals)
-        del vals
-
-        group["step"] += 1
-
-        stochastic_lerp_(exp_avg_list, grad_list, beta_debias(beta, group["step"]))
-
-        lr = -warmup(lr, group['step'], group['warmup_steps'])
-        lr = torch.empty((), dtype=torch.float32, device=grad_list[0].device).fill_(lr)
-
-        Q_list, exp_avg_list = list(Q_list), list(exp_avg_list)
-        for i, (p, g) in enumerate(zip(p_list, grad_list)):
-            q_orig = Q_list.pop(0)
-            ea = exp_avg_list.pop(0)
-            q = line_to_triu(q_orig) if store_triu_as_line else q_orig
-            _compilable_psgd_precond_grad_(q, self.state_(p)["exprs"][-1], ea, p, lr, weight_decay, self.clip_fn, group['caution'],
-                                           g)
-            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)

heavyball/foreach_adamw.py

@@ -1,63 +0,0 @@
-import torch
-import torch.optim
-from heavyball.utils import copy_stochastic_list_
-
-from .utils import warmup, exp_avg_sq_, beta_debias, update_param_, StatefulOptimizer, promote, decorator_knowngood
-
-
-@decorator_knowngood
-def _compilable_step_(y, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay, caution):
-    g32, exp_avg32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg, exp_avg_sq]]
-
-    torch._foreach_lerp_(exp_avg32, g32, 1 - beta_debias(beta1, step + 1))
-    denom = list(exp_avg_sq_(exp_avg_sq32, g32, beta_debias(beta2, step + 1), eps))
-
-    update_param_(y, exp_avg32, lr, decay, lambda p, e, l: p.addcdiv_(e, denom.pop(0), value=l), caution=caution,
-                  grad=g32)
-
-    copy_stochastic_list_(exp_avg, exp_avg32)
-    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
-
-
-class ForeachAdamW(StatefulOptimizer):
-    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):
-        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, storage_dtype=storage_dtype, mars=mars, caution=caution,
-                        mars_gamma=mars_gamma)
-        super().__init__(params, defaults, foreach)
-
-    def _step(self, group):
-        eps = group['eps']
-        decay = group['weight_decay']
-        k = group['k']
-
-        if not group['train_mode']:
-            raise Exception("Not in train mode!")
-
-        active_p = [p for p in group['params'] if p.grad is not None]
-
-        if not active_p:
-            return
-
-        storage_dtype = getattr(torch, group['storage_dtype'])
-
-        for p in active_p:
-            if 'exp_avg' not in self.state_(p):
-                self.state_(p)['exp_avg'] = torch.zeros_like(p.data, dtype=storage_dtype, memory_format=torch.preserve_format)
-                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=storage_dtype, memory_format=torch.preserve_format)
-
-        y, grad, exp_avg_sq, exp_avg = zip(
-            *[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
-
-        if group['mars']:
-            self.mars_correct_list(group, y, grad, group['mars_gamma'], group['betas'][0])
-
-        lr = -warmup(group['lr'], k + 1, group['warmup_steps'])
-        lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
-        step = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k)
-        _compilable_step_(y, grad, exp_avg_sq, exp_avg, group['betas'][0], group['betas'][1], step, lr, eps, decay,
-                          group['caution'])
-
-        group['k'] = k + 1

heavyball/foreach_adopt.py

@@ -1,83 +0,0 @@
-import torch
-import torch.optim
-from heavyball.utils import copy_stochastic_list_
-
-from .utils import warmup, beta_debias, update_param_, StatefulOptimizer, promote, decorator_knowngood
-
-
-@decorator_knowngood
-def _compilable_step_(y, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay, caution):
-    g32, exp_avg32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg, exp_avg_sq]]
-    update_param_(y, exp_avg, lr, decay, caution=caution, grad=g32)
-
-    beta1 = beta_debias(beta1, step)
-    denom = torch._foreach_sqrt(exp_avg_sq32)
-    torch._foreach_maximum_(denom, eps)
-    torch._foreach_mul_(exp_avg32, beta1)
-    [ea32.addcdiv_(g, d, value=1 - beta1) for ea32, g, d in zip(exp_avg32, g32, denom)]
-
-    beta2 = beta_debias(beta2, step + 1)
-    torch._foreach_mul_(exp_avg_sq32, beta2)
-    [eas32.addcmul_(g, g, value=1 - beta2) for eas32, g in zip(exp_avg_sq32, g32)]
-
-    copy_stochastic_list_(exp_avg, exp_avg32)
-    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
-
-
-class ForeachADOPT(StatefulOptimizer):
-
-    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):
-        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, storage_dtype=storage_dtype, mars=mars, caution=caution,
-                        mars_gamma=mars_gamma)
-        super().__init__(params, defaults, foreach)
-
-    def _step(self, group):
-        eps = group['eps']
-        decay = group['weight_decay']
-        k = group['k']
-
-        if not group['train_mode']:
-            raise Exception("Not in train mode!")
-
-        active_p = [p for p in group['params'] if p.grad is not None]
-
-        if not active_p:
-            return
-
-        storage_dtype = getattr(torch, group['storage_dtype'])
-
-        for p in active_p:
-            if 'exp_avg' not in self.state_(p):
-                self.state_(p)['exp_avg'] = torch.zeros_like(p.data, dtype=storage_dtype, memory_format=torch.preserve_format)
-                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=storage_dtype, memory_format=torch.preserve_format)
-
-        y, grad, exp_avg_sq, exp_avg = zip(
-            *[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
-
-        group['k'] = k + 1
-
-        if group['mars']:
-            self.mars_correct_list(group, y, grad, group['mars_gamma'], group['betas'][0])
-
-        if k > 1:
-            lr = -warmup(group['lr'], k - 1, group['warmup_steps'])
-            lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
-            k = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k)
-            _compilable_step_(y, grad, exp_avg_sq, exp_avg, group['betas'][0], group['betas'][1], k, lr, eps, decay, group['caution'])
-            return
-
-        grad = [promote(g) for g in grad]
-        if k > 0:
-            beta1 = beta_debias(group['betas'][0], k)
-            denom = torch._foreach_sqrt(exp_avg_sq)
-            torch._foreach_maximum_(denom, eps)
-            torch._foreach_mul_(exp_avg, beta1)
-            torch._foreach_addcdiv_(exp_avg, grad, denom, 1 - beta1)
-
-        beta2 = beta_debias(group['betas'][1], k + 1)
-        torch._foreach_mul_(exp_avg_sq, beta2)
-        torch._foreach_addcmul_(exp_avg_sq, grad, grad, value=1 - beta2)
-        del grad

heavyball/foreach_laprop.py

@@ -1,67 +0,0 @@
-import torch
-import torch.optim
-
-from .utils import warmup, exp_avg_sq_, beta_debias, update_param_, StatefulOptimizer, promote, copy_stochastic_list_, decorator_knowngood
-
-
-@decorator_knowngood
-def _compilable_step_(y, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay, caution):
-    g32, exp_avg32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg, exp_avg_sq]]
-
-    denom = exp_avg_sq_(exp_avg_sq32, g32, beta_debias(beta2, step), eps)
-
-    beta1 = beta_debias(beta1, step)
-    torch._foreach_mul_(exp_avg32, beta1)
-    [ea32.addcdiv_(g, d, value=1 - beta1) for ea32, g, d in zip(exp_avg32, g32, denom)]
-
-    update_param_(y, exp_avg32, lr, decay, caution=caution, grad=g32)
-
-    copy_stochastic_list_(exp_avg, exp_avg32)
-    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
-
-
-class ForeachLaProp(StatefulOptimizer):
-
-    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=1,
-                 foreach: bool = True, storage_dtype: str = 'float32', mars: bool = False, caution: bool = False,
-                 mars_gamma: float = 0.0025):
-        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, storage_dtype=storage_dtype, mars=mars, caution=caution,
-                        mars_gamma=mars_gamma)
-        super().__init__(params, defaults, foreach)
-
-    def _step(self, group):
-        eps = group['eps']
-        decay = group['weight_decay']
-        k = group['k']
-
-        if not group['train_mode']:
-            raise Exception("Not in train mode!")
-
-        active_p = [p for p in group['params'] if p.grad is not None]
-
-        if not active_p:
-            return
-
-        storage_dtype = getattr(torch, group['storage_dtype'])
-
-        for p in active_p:
-            if 'exp_avg' not in self.state_(p):
-                self.state_(p)['exp_avg'] = torch.zeros_like(p.data, dtype=storage_dtype, memory_format=torch.preserve_format)
-                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=storage_dtype, memory_format=torch.preserve_format)
-
-        y, grad, exp_avg_sq, exp_avg = zip(
-            *[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg'])  #
-              for p in active_p])
-
-        if group['mars']:
-            self.mars_correct_list(group, y, grad, group['mars_gamma'], group['betas'][0])
-
-        lr = -warmup(group['lr'], k + 1, group['warmup_steps'])
-        lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
-        step = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k + 1)
-
-        _compilable_step_(y, grad, exp_avg_sq, exp_avg, group['betas'][0], group['betas'][1], step, lr, eps, decay,
-                          group['caution'])
-
-        group['k'] = k + 1