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

heavyball/__init__.py

@@ -1,3 +1,5 @@
+from .cached_psgd_kron import ForeachCachedPSGDKron
+from .delayed_psgd import ForeachDelayedPSGD
 from .foreach_adamw import ForeachAdamW
 from .foreach_adopt import ForeachADOPT
 from .foreach_laprop import ForeachLaProp
@@ -12,11 +14,31 @@ from .precond_schedule_sfpsoap import PrecondScheduleSFPaLMSOAP
 from .psgd_kron import ForeachPSGDKron
 from .pure_psgd import ForeachPurePSGD
 from .schedule_free_palm_foreach_soap import SFPaLMForeachSOAP
-from .delayed_psgd import ForeachDelayedPSGD
 
 PalmForEachSoap = PaLMForeachSOAP
 
+PaLMSOAP = PaLMForeachSOAP
+PaLMSFAdamW = PaLMForeachSFAdamW
+PaLMSFSoap = SFPaLMForeachSOAP
+PaLMForeachSOAP = PaLMForeachSOAP
+PrecondScheduleSFPaLMSOAP = PrecondScheduleSFPaLMSOAP
+SOAP = ForeachSOAP
+SFAdamW = ForeachSFAdamW
+LaProp = ForeachLaProp
+ADOPT = ForeachADOPT
+PrecondScheduleForeachSOAP = PrecondScheduleForeachSOAP
+PrecondSchedulePaLMForeachSOAP = PrecondSchedulePaLMForeachSOAP
+PSGDKron = ForeachPSGDKron
+AdamW = ForeachAdamW
+PurePSGD = ForeachPurePSGD
+PaLMPAdam = ForeachPaLMPAdam
+DelayedPSGD = ForeachDelayedPSGD
+CachedPSGDKron = ForeachCachedPSGDKron
+
 __all__ = ['PalmForEachSoap', 'PaLMForeachSFAdamW', 'PaLMForeachSOAP', 'SFPaLMForeachSOAP', 'PrecondScheduleSFPaLMSOAP',
            'ForeachSOAP', 'ForeachSFAdamW', 'ForeachLaProp', 'ForeachADOPT', 'PrecondScheduleForeachSOAP',
-           'PrecondSchedulePaLMForeachSOAP', 'ForeachPSGDKron', 'ForeachAdamW', 'ForeachPurePSGD',
-           'ForeachPaLMPAdam', 'ForeachDelayedPSGD']
+           'PrecondSchedulePaLMForeachSOAP', 'ForeachPSGDKron', 'ForeachAdamW', 'ForeachPurePSGD', 'ForeachPaLMPAdam',
+           'ForeachDelayedPSGD', 'ForeachCachedPSGDKron',  #
+           'PaLMSOAP', 'PaLMSFAdamW', 'PaLMSFSoap', 'PaLMSFAdamW', 'PaLMForeachSOAP', 'PrecondScheduleSFPaLMSOAP',
+           'SOAP', 'SFAdamW', 'LaProp', 'ADOPT', 'PSGDKron', 'AdamW', 'PurePSGD', 'PaLMPAdam', 'DelayedPSGD',
+           'CachedPSGDKron']

heavyball/cached_psgd_kron.py

@@ -0,0 +1,141 @@
+"""
+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 einsum_base
+
+from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust_region_clip_, PSGDBase, \
+    precond_update_prob_schedule, split_p_and_g_in_group, line_to_triu, triu_to_line, set_, einsum_base
+
+
+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.
+        b1 (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):
+        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 preconditioner_update_probability is None:
+            preconditioner_update_probability = precond_update_prob_schedule()
+        if clip_fn is None:
+            clip_fn = lambda x: trust_region_clip_(x, 0.9, 1.5)
+        self.preconditioner_update_probability = preconditioner_update_probability
+        self.clip_fn = clip_fn
+
+        defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
+                        min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
+                        momentum_into_precond_update=momentum_into_precond_update, precond_lr=0.1,
+                        # precond lr hardcoded to 0.1
+                        precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
+                        step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
+                        store_triu_as_line=store_triu_as_line)
+        super().__init__(params, defaults)
+
+        self._prob_step = 0
+
+    def _step(self, group):
+        # update preconditioners all together
+        update_prob = self.preconditioner_update_probability
+        if callable(update_prob):
+            update_prob = update_prob(self._prob_step)
+        do_update = self.rng.random() < update_prob
+        self._prob_step += 1
+
+        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']
+
+        vals = []
+
+        for p, g in split_p_and_g_in_group(group):
+            state = self.state_(p)
+
+            if 'Q' not in state:
+                state["exp_avg"] = 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)
+                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
+
+        torch._foreach_lerp_(exp_avg_list, grad_list, (1 - beta) / (1 - beta ** group["step"]))
+
+        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 do_update:
+                q = line_to_triu(q_orig) if store_triu_as_line else q_orig
+                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)
+                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_)
+
+            set_(g, torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea))
+        grad_list = self.clip_fn(grad_list)
+
+        lr = -warmup(lr, group['step'], group['warmup_steps'])
+        update_param_(p_list, grad_list, lr, weight_decay)

heavyball/delayed_psgd.py

@@ -5,8 +5,8 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 """
 
 import torch
-
 from heavyball.utils import copy_stochastic_list_
+
 from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust_region_clip_, PSGDBase, \
     precond_update_prob_schedule, split_p_and_g_in_group, triu_to_line, line_to_triu, set_
 
@@ -38,7 +38,7 @@ 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):
+                 split: bool = False, clip_fn: callable = None, store_triu_as_line: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -58,18 +58,13 @@ class ForeachDelayedPSGD(PSGDBase):
                         momentum_into_precond_update=momentum_into_precond_update, precond_lr=0.1,
                         # precond lr hardcoded to 0.1
                         precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
-                        step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split)
+                        step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
+                        store_triu_as_line=store_triu_as_line)
         super().__init__(params, defaults)
 
         self._prob_step = 0
 
-    @torch.no_grad()
-    def step(self, closure=None):
-        loss = None
-        if closure is not None:
-            with torch.enable_grad():
-                loss = closure()
-
+    def _step(self, group):
         # update preconditioners all together
         update_prob = self.preconditioner_update_probability
         if callable(update_prob):
@@ -77,55 +72,52 @@ class ForeachDelayedPSGD(PSGDBase):
         do_update = self.rng.random() < update_prob
         self._prob_step += 1
 
-        for group in self.param_groups:
-            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']
-
-            vals = []
-
-            for p, g in split_p_and_g_in_group(group):
-                state = self.state_(p)
+        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']
 
-                if 'Q' not in state:
-                    state["exp_avg"] = 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)
-                    state["Q"] = triu_to_line(Q)
+        vals = []
 
-                vals.append((p, g, state["exp_avg"], state["Q"]))
+        for p, g in split_p_and_g_in_group(group):
+            state = self.state_(p)
 
-            if not vals:
-                continue
+            if 'Q' not in state:
+                state["exp_avg"] = 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)
+                state["Q"] = triu_to_line(Q) if store_triu_as_line else Q
 
-            p_list, grad_list, exp_avg_list, Q_list = zip(*vals)
-            del vals
+            vals.append((p, g, state["exp_avg"], state["Q"]))
 
-            group["step"] += 1
+        if not vals:
+            return
 
-            torch._foreach_lerp_(exp_avg_list, grad_list, (1 - beta) / (1 - beta ** group["step"]))
+        p_list, grad_list, exp_avg_list, Q_list = zip(*vals)
+        del vals
 
-            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)
-                self.balance(do_update, [g], [q])
-                new = psgd_precond_grad(q, self.state_(p)["exprs"], ea)
+        group["step"] += 1
 
-                if do_update:
-                    self.do_update([p], [ea if momentum_into_precond_update else g], [q], precond_lr, [q_orig])
-                set_(g, new)
+        torch._foreach_lerp_(exp_avg_list, grad_list, (1 - beta) / (1 - beta ** group["step"]))
 
-            grad_list = self.clip_fn(grad_list)
+        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
+            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.balance([g], [q])
+            set_(g, new)
 
-            lr = -warmup(lr, group['step'], group['warmup_steps'])
-            update_param_(p_list, grad_list, lr, weight_decay)
+        grad_list = self.clip_fn(grad_list)
 
-        return loss
+        lr = -warmup(lr, group['step'], group['warmup_steps'])
+        update_param_(p_list, grad_list, lr, weight_decay)

heavyball/foreach_adamw.py

@@ -10,42 +10,32 @@ class ForeachAdamW(StatefulOptimizer):
                         lr_max=-1.0, weight_decay=weight_decay)
         super().__init__(params, defaults)
 
-    def step(self, closure=None):
-        """Performs a single optimization step.
+    def _step(self, group):
+        eps = group['eps']
+        decay = group['weight_decay']
+        k = group['k']
 
-        Arguments:
-            closure (callable, optional): A closure that reevaluates the model
-                and returns the loss.
-        """
+        if not group['train_mode']:
+            raise Exception("Not in train mode!")
 
-        loss = None
-        if closure is not None:
-            loss = closure()
+        active_p = [p for p in group['params'] if p.grad is not None]
 
-        for group in self.param_groups:
-            eps = group['eps']
-            decay = group['weight_decay']
-            k = group['k']
+        if not active_p:
+            return
 
-            if not group['train_mode']:
-                raise Exception("Not in train mode!")
+        for p in active_p:
+            if 'exp_avg' not in self.state_(p):
+                self.state_(p)['exp_avg'] = torch.zeros_like(p.data, dtype=torch.float32)
+                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
 
-            active_p = [p for p in group['params'] if p.grad is not None]
+        y, grad, exp_avg_sq, exp_avg = zip(
+            *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
 
-            for p in active_p:
-                if 'exp_avg' not in self.state_(p):
-                    self.state_(p)['exp_avg'] = torch.zeros_like(p.data, dtype=torch.float32)
-                    self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
+        # Decay the first and second moment running average coefficient
+        torch._foreach_lerp_(exp_avg, grad, 1 - beta_debias(group['betas'][0], k + 1))
+        denom = exp_avg_sq_(exp_avg_sq, grad, beta_debias(group['betas'][1], k + 1), eps)
 
-            y, grad, exp_avg_sq, exp_avg = zip(
-                *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
-
-            # Decay the first and second moment running average coefficient
-            torch._foreach_lerp_(exp_avg, grad, 1 - beta_debias(group['betas'][0], k + 1))
-            denom = exp_avg_sq_(exp_avg_sq, grad, beta_debias(group['betas'][1], k + 1), eps)
-
-            # Normalize grad in-place for memory efficiency
-            lr = -warmup(group['lr'], k + 1, group['warmup_steps'])
-            update_param_(y, exp_avg, lr, decay, lambda p, e, l: torch._foreach_addcdiv_(p, e, denom, l))
-            group['k'] = k + 1
-        return loss
+        # Normalize grad in-place for memory efficiency
+        lr = -warmup(group['lr'], k + 1, group['warmup_steps'])
+        update_param_(y, exp_avg, lr, decay, lambda p, e, l: torch._foreach_addcdiv_(p, e, denom, l))
+        group['k'] = k + 1

heavyball/foreach_adopt.py

@@ -11,51 +11,41 @@ class ForeachADOPT(StatefulOptimizer):
                         lr_max=-1.0, weight_decay=weight_decay)
         super().__init__(params, defaults)
 
-    def step(self, closure=None):
-        """Performs a single optimization step.
-
-        Arguments:
-            closure (callable, optional): A closure that reevaluates the model
-                and returns the loss.
-        """
-
-        loss = None
-        if closure is not None:
-            loss = closure()
-
-        for group in self.param_groups:
-            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]
-
-            for p in active_p:
-                if 'exp_avg' not in self.state_(p):
-                    self.state_(p)['exp_avg'] = torch.zeros_like(p.data, dtype=torch.float32)
-                    self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
-
-            y, grad, exp_avg_sq, exp_avg = zip(
-                *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
-
-            if k > 1:
-                lr = -warmup(group['lr'], k - 1, group['warmup_steps'])
-
-                update_param_(y, exp_avg, lr, decay)
-            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
-
-            group['k'] = k + 1
-        return loss
+    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
+
+        for p in active_p:
+            if 'exp_avg' not in self.state_(p):
+                self.state_(p)['exp_avg'] = torch.zeros_like(p.data, dtype=torch.float32)
+                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
+
+        y, grad, exp_avg_sq, exp_avg = zip(
+            *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
+
+        if k > 1:
+            lr = -warmup(group['lr'], k - 1, group['warmup_steps'])
+
+            update_param_(y, exp_avg, lr, decay)
+        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
+
+        group['k'] = k + 1

heavyball/foreach_laprop.py

@@ -11,46 +11,36 @@ class ForeachLaProp(StatefulOptimizer):
                         lr_max=-1.0, weight_decay=weight_decay)
         super().__init__(params, defaults)
 
-    def step(self, closure=None):
-        """Performs a single optimization step.
+    def _step(self, group):
+        eps = group['eps']
+        decay = group['weight_decay']
+        k = group['k']
 
-        Arguments:
-            closure (callable, optional): A closure that reevaluates the model
-                and returns the loss.
-        """
+        if not group['train_mode']:
+            raise Exception("Not in train mode!")
 
-        loss = None
-        if closure is not None:
-            loss = closure()
+        active_p = [p for p in group['params'] if p.grad is not None]
 
-        for group in self.param_groups:
-            eps = group['eps']
-            decay = group['weight_decay']
-            k = group['k']
+        if not active_p:
+            return
 
-            if not group['train_mode']:
-                raise Exception("Not in train mode!")
+        for p in active_p:
+            if 'exp_avg' not in self.state_(p):
+                self.state_(p)['exp_avg'] = torch.zeros_like(p.data, dtype=torch.float32)
+                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
 
-            active_p = [p for p in group['params'] if p.grad is not None]
+        y, grad, exp_avg_sq, exp_avg = zip(
+            *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
 
-            for p in active_p:
-                if 'exp_avg' not in self.state_(p):
-                    self.state_(p)['exp_avg'] = torch.zeros_like(p.data, dtype=torch.float32)
-                    self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
+        # Decay the first and second moment running average coefficient
+        denom = exp_avg_sq_(exp_avg_sq, grad, beta_debias(group['betas'][1], k + 1), eps)
+        beta1 = beta_debias(group['betas'][0], k + 1)
+        torch._foreach_mul_(exp_avg, beta1)
+        torch._foreach_addcdiv_(exp_avg, grad, denom, 1 - beta1)
+        del grad
 
-            y, grad, exp_avg_sq, exp_avg = zip(
-                *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
+        # Normalize grad in-place for memory efficiency
+        lr = -warmup(group['lr'], k + 1, group['warmup_steps'])
+        update_param_(y, exp_avg, lr, decay)
 
-            # Decay the first and second moment running average coefficient
-            denom = exp_avg_sq_(exp_avg_sq, grad, beta_debias(group['betas'][1], k + 1), eps)
-            beta1 = beta_debias(group['betas'][0], k + 1)
-            torch._foreach_mul_(exp_avg, beta1)
-            torch._foreach_addcdiv_(exp_avg, grad, denom, 1 - beta1)
-            del grad
-
-            # Normalize grad in-place for memory efficiency
-            lr = -warmup(group['lr'], k + 1, group['warmup_steps'])
-            update_param_(y, exp_avg, lr, decay)
-
-            group['k'] = k + 1
-        return loss
+        group['k'] = k + 1

heavyball/foreach_sfadamw.py

@@ -13,52 +13,42 @@ class ForeachSFAdamW(ScheduleFree):
                         foreach=foreach)
         super().__init__(params, defaults)
 
-    def step(self, closure=None):
-        """Performs a single optimization step.
+    def _step(self, group):
+        eps = group['eps']
+        decay = group['weight_decay']
+        k = group['k']
 
-        Arguments:
-            closure (callable, optional): A closure that reevaluates the model
-                and returns the loss.
-        """
+        if not group['train_mode']:
+            raise Exception("Not in train mode!")
 
-        loss = None
-        if closure is not None:
-            loss = closure()
+        active_p = [p for p in group['params'] if p.grad is not None]
 
-        for group in self.param_groups:
-            eps = group['eps']
-            decay = group['weight_decay']
-            k = group['k']
+        if not active_p:
+            return
 
-            if not group['train_mode']:
-                raise Exception("Not in train mode!")
+        for p in active_p:
+            if 'z' not in self.state_(p):
+                self.state_(p)['z'] = torch.clone(p.data)
+                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
 
-            active_p = [p for p in group['params'] if p.grad is not None]
+        y, grad, exp_avg_sq, z = zip(
+            *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['z']) for p in active_p])
 
-            for p in active_p:
-                if 'z' not in self.state_(p):
-                    self.state_(p)['z'] = torch.clone(p.data)
-                    self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
+        # Decay the first moment running average coefficient
+        old_debiased = beta_debias(group['betas'][1], k + 1)
 
-            y, grad, exp_avg_sq, z = zip(
-                *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['z']) for p in active_p])
+        # Decay the first and second moment running average coefficient
+        denom = exp_avg_sq_(exp_avg_sq, grad, old_debiased, eps)
 
-            # Decay the first moment running average coefficient
-            old_debiased = beta_debias(group['betas'][1], k + 1)
+        # Normalize grad in-place for memory efficiency
+        torch._foreach_div_(grad, denom)
 
-            # Decay the first and second moment running average coefficient
-            denom = exp_avg_sq_(exp_avg_sq, grad, old_debiased, eps)
+        # Weight decay calculated at y
+        if decay != 0:
+            torch._foreach_add_(grad, y, alpha=decay)
 
-            # Normalize grad in-place for memory efficiency
-            torch._foreach_div_(grad, denom)
+        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)
 
-            # Weight decay calculated at y
-            if decay != 0:
-                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['k'] = k + 1
-        return loss
+        group['k'] = k + 1