heavyball 0.25.1__py3-none-any.whl → 1.1.0__py3-none-any.whl

heavyball/foreach_sfadamw.py

@@ -1,69 +0,0 @@
-import torch
-import torch.optim
-
-from .utils import get_ckp1, copy_stochastic_list_, warmup, ScheduleFree, exp_avg_sq_, beta_debias, promote, _compilable_schedule_free_, decorator_knowngood
-
-
-@decorator_knowngood
-def _compilable_step_(y, grad, exp_avg_sq, z, beta1, beta2, step, ckp1, eps, decay, lr):
-    old_debiased2 = beta_debias(beta2, step)
-
-    g32 = [promote(g_) for g_ in grad]
-    exp_avg_sq32 = [promote(e_) for e_ in exp_avg_sq]
-
-    denom = exp_avg_sq_(exp_avg_sq32, g32, old_debiased2, eps)
-    torch._foreach_div_(g32, denom)
-    if decay != 0:
-        torch._foreach_add_(g32, y, alpha=decay)
-    for p, z_, g in zip(y, z, g32):
-        _compilable_schedule_free_(p, z_, ckp1, g, lr, beta1)
-
-    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
-
-
-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: bool = True, storage_dtype: str = 'float32', mars: bool = False,
-                 caution: bool = False, mars_gamma: float = 0.0025):
-
-        assert not caution, "Caution not implemented for SFAdamW"
-
-        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, 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 'z' not in self.state_(p):
-                self.state_(p)['z'] = torch.clone(p.data, 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, z = zip(*[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['z'])  #
-                                       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'])
-        ckp1, group['weight_sum'] = get_ckp1(lr, group['weight_lr_power'], group['weight_sum'], group['r'], k + 1)
-
-        step = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k + 1)
-        ckp1 = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(ckp1)
-        lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
-        _compilable_step_(y, grad, exp_avg_sq, z, group['betas'][0], group['betas'][1], step, ckp1, eps, decay, lr)
-        group['k'] = k + 1

heavyball/foreach_soap.py

@@ -1,93 +0,0 @@
-import torch
-
-from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, set_, \
-    StatefulOptimizer, exp_avg_
-
-
-class ForeachSOAP(StatefulOptimizer):
-    """
-    ForeachSOAP
-
-    Sources:
-        Baseline SOAP:
-            SOAP: Improving and Stabilizing Shampoo using Adam
-            Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
-            https://arxiv.org/abs/2409.11321
-            https://github.com/nikhilvyas/SOAP
-
-        ScheduleFree:
-            The Road Less Scheduled
-            Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky
-            https://arxiv.org/abs/2405.15682
-            https://github.com/facebookresearch/schedule_free
-    """
-
-    def __init__(self, params, lr: float = 3e-3, betas=(0.9, 0.95), shampoo_beta: float = 0.95, eps: float = 1e-8,
-                 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, foreach: bool = True, mars: bool = False, caution: bool = False,
-                 mars_gamma: float = 0.0025):
-        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, 'mars': mars,
-                    'caution': caution, 'mars_gamma': mars_gamma}
-        super().__init__(params, defaults, foreach)
-        self._data_format = data_format
-
-    def _step(self, group):
-        vals = []
-        step = 0
-
-        max_precond_dim = group['max_precond_dim']
-        precondition_1d = group['precondition_1d']
-
-        for p, g in self.split_p_and_g_in_group(group, beta1=group['betas'][0]):
-            state = self.state_(p)
-            step = state['step'] = state.get("step", -1) + 1
-
-            if "exp_avg" not in state:
-                state["exp_avg"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
-                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
-                init_preconditioner(g, state, max_precond_dim, precondition_1d)
-                update_preconditioner(g, state, max_precond_dim, precondition_1d, 0, True)
-                continue  # first step is skipped so that we never use the current gradients in the projection.
-
-            # Projecting gradients to the eigenbases of Shampoo's preconditioner
-            # i.e. projecting to the eigenbases of matrices in state['GG']
-            grad_projected = project(g, state['Q'], False)
-            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
-            vals.append((p, g, grad_projected, exp_avg, exp_avg_sq))
-
-        if not vals:
-            return
-
-        p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
-        beta1, beta2 = group["betas"]
-
-        old_debiased2 = beta_debias(beta2, step)
-
-        # Decay the first and second moment running average coefficient
-        # In-place operations to update the averages at the same time
-        step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
-
-        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
-
-        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
-            state = self.state_(p)
-
-            d = exp_avg_(ea, eas, g, gp, beta1, beta2, step_tensor)[0]
-            # Projecting the exponential moving average of gradients to the eigenbases of Shampoo's preconditioner
-            # i.e. projecting to the eigenbases of matrices in state['GG']
-            exp_avg_projected = project(ea, state['Q'], False)
-
-            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
-            # to the original space
-            # CANT DO /= HERE AS EXP_AVG MAY POINT TO THE BUFFER
-            precond = project(exp_avg_projected / d, state['Q'], True)
-
-            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2,
-                                  step > 0 and step % group['precondition_frequency'] == 0)
-
-            update_param_([p], [precond], step_size, group["weight_decay"], caution=group['caution'], grad=[g])

heavyball/foreach_solp.py

@@ -1,89 +0,0 @@
-import torch
-
-from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, set_, \
-    StatefulOptimizer, laprop_exp_avg_
-
-
-class ForeachSOLP(StatefulOptimizer):
-    """
-    ForeachSOLP
-
-    Sources:
-        LaProp:
-            LaProp: Separating Momentum and Adaptivity in Adam
-            Liu Ziyin, Zhikang T.Wang, Masahito Ueda
-            https://arxiv.org/abs/2002.04839
-            https://github.com/ClashLuke/HeavyBall/blob/main/heavyball/foreach_laprop.py
-
-        Baseline SOAP:
-            SOAP: Improving and Stabilizing Shampoo using Adam
-            Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
-            https://arxiv.org/abs/2409.11321
-            https://github.com/nikhilvyas/SOAP
-
-        ScheduleFree:
-            The Road Less Scheduled
-            Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky
-            https://arxiv.org/abs/2405.15682
-            https://github.com/facebookresearch/schedule_free
-    """
-
-    def __init__(self, params, lr: float = 3e-3, betas=(0.9, 0.95), shampoo_beta: float = 0.95, eps: float = 1e-8,
-                 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, foreach: bool = True, mars: bool = False, caution: bool = False,
-                 mars_gamma: float = 0.0025):
-        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, 'mars': mars,
-                    'caution': caution, 'mars_gamma': mars_gamma}
-        super().__init__(params, defaults, foreach)
-        self._data_format = data_format
-
-    def _step(self, group):
-        vals = []
-        step = 0
-
-        max_precond_dim = group['max_precond_dim']
-        precondition_1d = group['precondition_1d']
-
-        for p, g in self.split_p_and_g_in_group(group, beta1=group['betas'][0]):
-            state = self.state_(p)
-            step = state['step'] = state.get("step", -1) + 1
-
-            if "exp_avg" not in state:
-                state["exp_avg"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
-                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
-                init_preconditioner(g, state, max_precond_dim, precondition_1d)
-                update_preconditioner(g, state, max_precond_dim, precondition_1d, 0, True)
-                continue  # first step is skipped so that we never use the current gradients in the projection.
-
-            # Projecting gradients to the eigenbases of Shampoo's preconditioner
-            # i.e. projecting to the eigenbases of matrices in state['GG']
-            grad_projected = project(g, state['Q'], False)
-            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
-            vals.append((p, g, grad_projected, exp_avg, exp_avg_sq))
-
-        if not vals:
-            return
-
-        p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
-        beta1, beta2 = group["betas"]
-
-        old_debiased2 = beta_debias(beta2, step)
-
-        # Decay the first and second moment running average coefficient
-        # In-place operations to update the averages at the same time
-        step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
-
-        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
-
-        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
-            laprop_exp_avg_(ea, eas, gp, beta1, beta2, step_tensor)
-            state = self.state_(p)
-            precond = project(ea, state['Q'], True)
-
-            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2, step > 0 and step % group['precondition_frequency'] == 0)
-            update_param_([p], [precond], step_size, group["weight_decay"], caution=group['caution'], grad=[g])

heavyball/p_adam.py

@@ -1,121 +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 triu_to_line, line_to_triu, identity, stochastic_lerp_
-
-from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, PSGDBase, exp_avg_sq_, beta_debias, \
-    promote
-
-
-class ForeachPaLMPAdam(PSGDBase):
-    """
-    Kronecker Factorized Adam with PSGD preconditioner
-
-    Args:
-        params (iterable): Iterable of parameters to optimize or dicts defining
-            parameter groups.
-        lr (float): Learning rate.
-        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, 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, 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, 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 <= weight_decay:
-            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
-        if betas[0] is not None:
-            beta = betas[0]
-
-        if clip_fn is None:
-            clip_fn = identity
-
-        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=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, storage_dtype=storage_dtype,
-                        mars=mars, caution=caution, mars_gamma=mars_gamma)
-        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']
-        memory_save_mode = group['memory_save_mode']
-        precond_lr = group['precond_lr']
-        weight_decay = group['weight_decay']
-        lr = group['lr']
-        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=group['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)
-                state['exp_avg_sq'] = 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["Q"], state['exp_avg'], state['exp_avg_sq']))
-
-        if not vals:
-            return
-
-        p_list, grad_list, Q_list, exp_avg, exp_avg_sq = zip(*vals)
-        del vals
-
-        group["step"] += 1
-
-        Q_triu = [line_to_triu(q) if store_triu_as_line else q for q in Q_list]
-        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)
-        stochastic_lerp_(exp_avg, grad_list, 1 - beta_debias(group['beta'], group['step']))
-
-        beta2 = 1 - group['step'] ** -group['beta2_scale']
-
-        lr = -warmup(lr, group['step'], group['warmup_steps'])
-
-        for p, Q, g, ea, eas in zip(p_list, Q_triu, grad_list, exp_avg, exp_avg_sq):
-            gc = g.clone() if group['caution'] else None
-            psgd_precond_grad(True, self.state_(p)["exprs"][-1], g, *Q)
-            ea = psgd_precond_grad(False, self.state_(p)["exprs"][-1], ea, *Q)
-            exp_avg_sq_(eas, g, beta_debias(beta2, group['step']), 1e-8, out=g)
-            torch.div(ea, g, out=g)
-            """
-            divide by g here, because g == denom (from exp_avg_sq_(out=g)), avoids denom allocation
-            divide into g so we can deallocate ea, avoids one allocation (-> less memory than equivalent foreach)
-            """
-            update_param_([p], self.clip_fn([g]), lr, weight_decay, caution=group['caution'], grad=gc)

heavyball/palm_foreach_sfadamw.py

@@ -1,77 +0,0 @@
-import torch
-import torch.optim
-
-from .utils import warmup, ScheduleFree, exp_avg_sq_, beta_debias, get_ckp1, promote, \
-    _compilable_schedule_free_, copy_stochastic_list_, decorator_knowngood
-
-
-@decorator_knowngood
-def _compilable_step_(y, grad, exp_avg_sq, z, beta1, beta2, step, ckp1, eps, decay, lr):
-    old_debiased2 = beta_debias(beta2, step)
-
-    g32 = [promote(g_) for g_ in grad]
-    exp_avg_sq32 = [promote(e_) for e_ in exp_avg_sq]
-
-    denom = exp_avg_sq_(exp_avg_sq32, g32, old_debiased2, eps)
-    torch._foreach_div_(g32, denom)
-    if decay != 0:
-        torch._foreach_add_(g32, y, alpha=decay)
-    for p, z_, g in zip(y, z, g32):
-        _compilable_schedule_free_(p, z_, ckp1, g, lr, beta1)
-
-    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
-
-
-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, foreach: bool = True, storage_dtype: str = 'float32',
-                 mars: bool = False, caution: bool = False, mars_gamma: float = 0.0025):
-        if betas[0] is not None:
-            beta = betas[0]
-
-        assert not caution, "Caution not implemented for SFAdamW"
-
-        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, 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 'z' not in self.state_(p):
-                self.state_(p)['z'] = torch.clone(p.data, memory_format=torch.preserve_format)
-                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=storage_dtype, memory_format=torch.preserve_format)
-
-        # Decay the first moment running average coefficient
-        beta2 = 1 - (k + 1) ** -group['beta2_scale']
-
-        y, grad, exp_avg_sq, z = zip(*[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['z'])  #
-                                       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'])
-        ckp1, group['weight_sum'] = get_ckp1(lr, group['weight_lr_power'], group['weight_sum'], group['r'], k + 1)
-
-        step = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k + 1)
-        ckp1 = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(ckp1)
-        beta2 = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(beta2)
-        lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
-        _compilable_step_(y, grad, exp_avg_sq, z, group['beta'], beta2, step, ckp1, eps, decay, lr)
-        group['k'] = k + 1

heavyball/palm_foreach_soap.py

@@ -1,101 +0,0 @@
-import torch
-
-from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, set_, \
-    StatefulOptimizer, exp_avg_
-
-
-class PaLMForeachSOAP(StatefulOptimizer):
-    """
-    PaLMForeachSOAP
-
-    Sources:
-        Baseline SOAP:
-            SOAP: Improving and Stabilizing Shampoo using Adam
-            Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
-            https://arxiv.org/abs/2409.11321
-            https://github.com/nikhilvyas/SOAP
-
-        ScheduleFree:
-            The Road Less Scheduled
-            Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky
-            https://arxiv.org/abs/2405.15682
-            https://github.com/facebookresearch/schedule_free
-
-        Beta2 Schedule:
-            PaLM: Scaling Language Modeling with Pathways
-            Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, Noah Fiedel
-            https://arxiv.org/abs/2204.02311
-    """
-
-    def __init__(self, params, lr: float = 3e-3, beta=0.9, betas=(None, None), shampoo_beta: float = 0.95,
-                 eps: float = 1e-8, 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,
-                 beta2_scale: float = 0.8, split: bool = False, foreach: bool = True, mars: bool = False,
-                 caution: bool = False, mars_gamma: float = 0.0025):
-        if betas[0] is not None:
-            beta = betas[0]
-        defaults = {"lr": lr, "beta": beta, "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, 'beta2_scale': beta2_scale,
-                    'split': split, 'mars': mars, 'caution': caution, 'mars_gamma': mars_gamma}
-        super().__init__(params, defaults, foreach)
-        self._data_format = data_format
-
-    def _step(self, group):
-        vals = []
-        step = 0
-
-        max_precond_dim = group['max_precond_dim']
-        precondition_1d = group['precondition_1d']
-
-        for p, g in self.split_p_and_g_in_group(group, beta1=group['beta']):
-            state = self.state_(p)
-            step = state['step'] = state.get("step", -1) + 1
-
-            if "exp_avg" not in state:
-                state["exp_avg"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
-                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
-                init_preconditioner(g, state, max_precond_dim, precondition_1d)
-                update_preconditioner(g, state, max_precond_dim, precondition_1d, 0, True)
-                continue  # first step is skipped so that we never use the current gradients in the projection.
-
-            # Projecting gradients to the eigenbases of Shampoo's preconditioner
-            # i.e. projecting to the eigenbases of matrices in state['GG']
-            grad_projected = project(g, state['Q'], False)
-            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
-            vals.append((p, g, grad_projected, exp_avg, exp_avg_sq))
-
-        if not vals:
-            return
-
-        p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
-        beta1 = group["beta"]
-
-        beta2 = 1 - step ** -group['beta2_scale']
-        old_debiased2 = beta_debias(beta2, step)
-
-        # Decay the first and second moment running average coefficient
-        # In-place operations to update the averages at the same time
-        beta2 = torch.empty((), dtype=torch.float32, device=p_list[0].device).fill_(beta2)
-        step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
-        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
-
-        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
-            state = self.state_(p)
-            d = exp_avg_(ea, eas, g, gp, beta1, beta2, step_tensor)[0]
-            # Projecting the exponential moving average of gradients to the eigenbases of Shampoo's preconditioner
-            # i.e. projecting to the eigenbases of matrices in state['GG']
-            exp_avg_projected = project(ea, state['Q'], False)
-
-            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
-            # to the original space
-            # CANT DO /= HERE AS EXP_AVG MAY POINT TO THE BUFFER
-            precond = project(exp_avg_projected / d, state['Q'], True)
-
-            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2,
-                                  step > 0 and step % group['precondition_frequency'] == 0)
-
-            update_param_([p], [precond], step_size, group["weight_decay"], caution=group['caution'], grad=[g])

heavyball/palm_foreach_solp.py

@@ -1,98 +0,0 @@
-import torch
-
-from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, set_, \
-    StatefulOptimizer, laprop_exp_avg_
-
-
-class PaLMForeachSOLP(StatefulOptimizer):
-    """
-    PaLMForeachSOAP
-
-    Sources:
-        LaProp:
-            LaProp: Separating Momentum and Adaptivity in Adam
-            Liu Ziyin, Zhikang T.Wang, Masahito Ueda
-            https://arxiv.org/abs/2002.04839
-            https://github.com/ClashLuke/HeavyBall/blob/main/heavyball/foreach_laprop.py
-
-        Baseline SOAP:
-            SOAP: Improving and Stabilizing Shampoo using Adam
-            Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
-            https://arxiv.org/abs/2409.11321
-            https://github.com/nikhilvyas/SOAP
-
-        ScheduleFree:
-            The Road Less Scheduled
-            Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky
-            https://arxiv.org/abs/2405.15682
-            https://github.com/facebookresearch/schedule_free
-
-        Beta2 Schedule:
-            PaLM: Scaling Language Modeling with Pathways
-            Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, Noah Fiedel
-            https://arxiv.org/abs/2204.02311
-    """
-
-    def __init__(self, params, lr: float = 3e-3, beta=0.9, betas=(None, None), shampoo_beta: float = 0.95,
-                 eps: float = 1e-8, 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,
-                 beta2_scale: float = 0.8, split: bool = False, foreach: bool = True, mars: bool = False,
-                 caution: bool = False, mars_gamma: float = 0.0025):
-        if betas[0] is not None:
-            beta = betas[0]
-        defaults = {"lr": lr, "beta": beta, "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, 'beta2_scale': beta2_scale,
-                    'split': split, 'mars': mars, 'caution': caution, 'mars_gamma': mars_gamma}
-        super().__init__(params, defaults, foreach)
-        self._data_format = data_format
-
-    def _step(self, group):
-        vals = []
-        step = 0
-
-        max_precond_dim = group['max_precond_dim']
-        precondition_1d = group['precondition_1d']
-
-        for p, g in self.split_p_and_g_in_group(group, beta1=group['beta']):
-            state = self.state_(p)
-            step = state['step'] = state.get("step", -1) + 1
-
-            if "exp_avg" not in state:
-                state["exp_avg"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
-                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
-                init_preconditioner(g, state, max_precond_dim, precondition_1d)
-                update_preconditioner(g, state, max_precond_dim, precondition_1d, 0, True)
-                continue  # first step is skipped so that we never use the current gradients in the projection.
-
-            # Projecting gradients to the eigenbases of Shampoo's preconditioner
-            # i.e. projecting to the eigenbases of matrices in state['GG']
-            grad_projected = project(g, state['Q'], False)
-            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
-            vals.append((p, g, grad_projected, exp_avg, exp_avg_sq))
-
-        if not vals:
-            return
-
-        p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
-        beta1 = group["beta"]
-
-        beta2 = 1 - step ** -group['beta2_scale']
-        old_debiased2 = beta_debias(beta2, step)
-
-        # Decay the first and second moment running average coefficient
-        # In-place operations to update the averages at the same time
-        beta2 = torch.empty((), dtype=torch.float32, device=p_list[0].device).fill_(beta2)
-        step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
-        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
-
-        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
-            laprop_exp_avg_(ea, eas, gp, beta1, beta2, step_tensor)
-            state = self.state_(p)
-            precond = project(ea, state['Q'], True)
-
-            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2, step > 0 and step % group['precondition_frequency'] == 0)
-            update_param_([p], [precond], step_size, group["weight_decay"], caution=group['caution'], grad=[g])