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

heavyball/foreach_soap.py

@@ -34,73 +34,59 @@ class ForeachSOAP(StatefulOptimizer):
         super().__init__(params, defaults)
         self._data_format = data_format
 
-    @torch.no_grad()
-    def step(self, closure=None):
-        """
-        Performs a single optimization step.
-
-        Arguments:
-            closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss.
-        """
-        if closure is None:
-            loss = None
-        else:
-            loss = closure()
-
-        for group in self.param_groups:
-            vals = []
-            step = 0
-
-            max_precond_dim = group['max_precond_dim']
-            precondition_1d = group['precondition_1d']
-
-            for p, g in split_p_and_g_in_group(group):
-                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)
-                    state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32)
-                    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:
-                continue
-
-            p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
-            beta1, beta2 = group["betas"]
-
-            old_debiased1 = beta_debias(beta1, step)
-            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
-            torch._foreach_mul_(exp_avg, old_debiased1)
-            torch._foreach_add_(exp_avg, grad, alpha=1 - old_debiased1)
-            denom = exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, group['eps'])
-
-            for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
-                state = self.state_(p)
-                # 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
-                set_(d, 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)
-
-            # Why does this have to be rebiased here?
-            step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
-            update_param_(p_list, denom, step_size, group["weight_decay"])
-        return loss
+    def _step(self, group):
+        vals = []
+        step = 0
+
+        max_precond_dim = group['max_precond_dim']
+        precondition_1d = group['precondition_1d']
+
+        for p, g in split_p_and_g_in_group(group):
+            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)
+                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32)
+                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_debiased1 = beta_debias(beta1, step)
+        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
+        torch._foreach_mul_(exp_avg, old_debiased1)
+        torch._foreach_add_(exp_avg, grad, alpha=1 - old_debiased1)
+        denom = exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, group['eps'])
+
+        for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
+            state = self.state_(p)
+            # 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
+            set_(d, 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)
+
+        # Why does this have to be rebiased here?
+        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
+        update_param_(p_list, denom, step_size, group["weight_decay"])

heavyball/p_adam.py

@@ -5,6 +5,7 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 """
 
 import torch
+from heavyball.utils import triu_to_line, line_to_triu
 
 from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, PSGDBase, precond_update_prob_schedule, \
     exp_avg_sq_, beta_debias, split_p_and_g_in_group
@@ -36,7 +37,8 @@ class ForeachPaLMPAdam(PSGDBase):
     def __init__(self, params, lr=0.001, weight_decay=0.0, preconditioner_update_probability=None,
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, betas=(None, None), beta: float = 0.9,
-                 beta2_scale: float = 0.8, merge_dims: bool = False, split: bool = False, clip_fn: callable = None):
+                 beta2_scale: float = 0.8, merge_dims: bool = False, 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 <= weight_decay:
@@ -57,18 +59,12 @@ class ForeachPaLMPAdam(PSGDBase):
                         # precond lr hardcoded to 0.1
                         precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
                         step=0, warmup_steps=warmup_steps, beta=beta, beta2_scale=beta2_scale, merge_dims=merge_dims,
-                        split=split)
+                        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):
@@ -76,57 +72,57 @@ class ForeachPaLMPAdam(PSGDBase):
         do_update = self.rng.random() < update_prob
         self._prob_step += 1
 
-        for group in self.param_groups:
-            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']
-
-            vals = []
+        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']
 
-            for p, g in split_p_and_g_in_group(group):
-                state = self.state_(p)
+        vals = []
 
-                if 'Q' not in state:
-                    state['exp_avg'] = torch.zeros_like(g)
-                    state['exp_avg_sq'] = torch.zeros_like(g)
-                    state["Q"], state["exprs"] = init_Q_exprs(p, precond_init_scale, max_size_triangular,
-                                                              min_ndim_triangular, memory_save_mode, dtype=g.dtype)
+        for p, g in split_p_and_g_in_group(group):
+            state = self.state_(p)
 
-                vals.append((p, g, state["Q"], state['exp_avg'], state['exp_avg_sq']))
+            if 'Q' not in state:
+                state['exp_avg'] = torch.zeros_like(g)
+                state['exp_avg_sq'] = torch.zeros_like(g)
+                Q, state["exprs"] = init_Q_exprs(p, precond_init_scale, max_size_triangular,
+                                                          min_ndim_triangular, memory_save_mode, dtype=g.dtype)
+                state['Q'] = triu_to_line(Q) if store_triu_as_line else Q
 
-            if not vals:
-                continue
+            vals.append((p, g, state["Q"], state['exp_avg'], state['exp_avg_sq']))
 
-            p_list, grad_list, Q_list, exp_avg, exp_avg_sq = zip(*vals)
-            del vals
+        if not vals:
+            return
 
-            group["step"] += 1
+        p_list, grad_list, Q_list, exp_avg, exp_avg_sq = zip(*vals)
+        del vals
 
-            self.balance(do_update, grad_list, Q_list)
-            if do_update:
-                self.do_update(p_list, grad_list, Q_list, precond_lr)
+        group["step"] += 1
 
-            torch._foreach_lerp_(exp_avg, grad_list, 1 - beta_debias(group['beta'], group['step']))
+        Q_triu = [line_to_triu(q) if store_triu_as_line else q for q in Q_list]
+        if do_update:
+            self.balance(grad_list, Q_triu)
+            self.do_update(p_list, grad_list, Q_triu, precond_lr, Q_list if store_triu_as_line else None)
 
-            beta2 = 1 - group['step'] ** -group['beta2_scale']
+        torch._foreach_lerp_(exp_avg, grad_list, 1 - beta_debias(group['beta'], group['step']))
 
-            for p, Q, g, ea, eas in zip(p_list, Q_list, grad_list, exp_avg, exp_avg_sq):
-                psgd_precond_grad(Q, self.state_(p)["exprs"], g, inplace=True)
-                ea = psgd_precond_grad(Q, self.state_(p)["exprs"], ea)
-                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)
-                """
+        beta2 = 1 - group['step'] ** -group['beta2_scale']
 
-            grad_list = self.clip_fn(grad_list)
+        for p, Q, g, ea, eas in zip(p_list, Q_triu, grad_list, exp_avg, exp_avg_sq):
+            psgd_precond_grad(Q, self.state_(p)["exprs"], g, inplace=True)
+            ea = psgd_precond_grad(Q, self.state_(p)["exprs"], ea)
+            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)
+            """
 
-            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/palm_foreach_sfadamw.py

@@ -9,58 +9,48 @@ class PaLMForeachSFAdamW(ScheduleFree):
                  weight_lr_power=2.0, beta2_scale: float = 0.8):
         if betas[0] is not None:
             beta = betas[0]
-        defaults = dict(lr=lr, beta=beta, eps=eps, r=r, k=0, warmup_steps=warmup_steps, train_mode=True,
-                        weight_sum=0.0, lr_max=-1.0, weight_lr_power=weight_lr_power, weight_decay=weight_decay,
+        defaults = dict(lr=lr, beta=beta, eps=eps, r=r, k=0, warmup_steps=warmup_steps, train_mode=True, weight_sum=0.0,
+                        lr_max=-1.0, weight_lr_power=weight_lr_power, weight_decay=weight_decay,
                         beta2_scale=beta2_scale)
         super().__init__(params, defaults)
 
-    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
+        beta2 = 1 - (k + 1) ** -group['beta2_scale']
+        old_debiased = beta_debias(beta2, 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
-            beta2 = 1 - (k + 1) ** -group['beta2_scale']
-            old_debiased = beta_debias(beta2, 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['beta'], 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['beta'],
-                                                 y, z, grad, group['r'], k + 1)
-
-            group['k'] = k + 1
-        return loss
+        group['k'] = k + 1

heavyball/palm_foreach_soap.py

@@ -43,73 +43,59 @@ class PaLMForeachSOAP(StatefulOptimizer):
         super().__init__(params, defaults)
         self._data_format = data_format
 
-    @torch.no_grad()
-    def step(self, closure=None):
-        """
-        Performs a single optimization step.
-
-        Arguments:
-            closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss.
-        """
-        if closure is None:
-            loss = None
-        else:
-            loss = closure()
-
-        for group in self.param_groups:
-            vals = []
-            step = 0
-
-            max_precond_dim = group['max_precond_dim']
-            precondition_1d = group['precondition_1d']
-
-            for p, g in split_p_and_g_in_group(group):
-                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)
-                    state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32)
-                    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:
-                continue
-
-            p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
-            beta1 = group["beta"]
-
-            beta2 = 1 - step ** -group['beta2_scale']
-            old_debiased1 = beta_debias(beta1, step)
-            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
-            torch._foreach_lerp_(exp_avg, grad, 1 - old_debiased1)
-            denom = exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, group['eps'])
-
-            for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
-                state = self.state_(p)
-                # 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
-                set_(d, 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)
-
-            # Why does this have to be rebiased here?
-            step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
-            update_param_(p_list, denom, step_size, group["weight_decay"])
-        return loss
+    def _step(self, group):
+        vals = []
+        step = 0
+
+        max_precond_dim = group['max_precond_dim']
+        precondition_1d = group['precondition_1d']
+
+        for p, g in split_p_and_g_in_group(group):
+            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)
+                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32)
+                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_debiased1 = beta_debias(beta1, step)
+        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
+        torch._foreach_lerp_(exp_avg, grad, 1 - old_debiased1)
+        denom = exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, group['eps'])
+
+        for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
+            state = self.state_(p)
+            # 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
+            set_(d, 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)
+
+        # Why does this have to be rebiased here?
+        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
+        update_param_(p_list, denom, step_size, group["weight_decay"])

heavyball/precond_schedule_foreach_soap.py

@@ -37,74 +37,60 @@ class PrecondScheduleForeachSOAP(StatefulOptimizer):
         self._data_format = data_format
         self.rng = random.Random(0x120983109)
 
-    @torch.no_grad()
-    def step(self, closure=None):
-        """
-        Performs a single optimization step.
-
-        Arguments:
-            closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss.
-        """
-        if closure is None:
-            loss = None
-        else:
-            loss = closure()
-
-        for group in self.param_groups:
-            vals = []
-            step = 0
-
-            max_precond_dim = group['max_precond_dim']
-            precondition_1d = group['precondition_1d']
-
-            for p, g in split_p_and_g_in_group(group):
-                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)
-                    state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32)
-                    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:
-                continue
-
-            p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
-            beta1, beta2 = group["betas"]
-
-            old_debiased1 = beta_debias(beta1, step)
-            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
-            torch._foreach_mul_(exp_avg, old_debiased1)
-            torch._foreach_add_(exp_avg, grad, alpha=1 - old_debiased1)
-            denom = exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, group['eps'])
-
-            update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
-            for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
-                state = self.state_(p)
-                # 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
-                set_(d, project(exp_avg_projected / d, state['Q'], True))
-
-                update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2,
-                                      update_precond)
-
-            # Why does this have to be rebiased here?
-            step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
-            update_param_(p_list, denom, step_size, group["weight_decay"])
-        return loss
+    def _step(self, group):
+        vals = []
+        step = 0
+
+        max_precond_dim = group['max_precond_dim']
+        precondition_1d = group['precondition_1d']
+
+        for p, g in split_p_and_g_in_group(group):
+            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)
+                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32)
+                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_debiased1 = beta_debias(beta1, step)
+        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
+        torch._foreach_mul_(exp_avg, old_debiased1)
+        torch._foreach_add_(exp_avg, grad, alpha=1 - old_debiased1)
+        denom = exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, group['eps'])
+
+        update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
+        for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
+            state = self.state_(p)
+            # 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
+            set_(d, project(exp_avg_projected / d, state['Q'], True))
+
+            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2,
+                                  update_precond)
+
+        # Why does this have to be rebiased here?
+        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
+        update_param_(p_list, denom, step_size, group["weight_decay"])