heavyball 0.19.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

heavyball/cached_delayed_psgd_kron.py

@@ -9,7 +9,7 @@ from typing import Optional
 import torch
 
 from .utils import update_param_, warmup, init_Q_exprs, trust_region_clip_, PSGDBase, split_p_and_g_in_group, \
-    line_to_triu, triu_to_line, set_, einsum_base, promote
+    line_to_triu, triu_to_line, einsum_base, promote, stochastic_lerp_, beta_debias
 
 
 class ForeachCachedDelayedPSGDKron(PSGDBase):
@@ -41,7 +41,8 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,  #
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,
+                 storage_dtype: str = 'float32',  #
                  # expert parameters
                  precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
@@ -58,7 +59,7 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
                         momentum_into_precond_update=momentum_into_precond_update, precond_lr=precond_lr,
                         precond_init_scale=precond_init_scale, step=0, warmup_steps=warmup_steps, merge_dims=merge_dims,
-                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
+                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype, storage_dtype=storage_dtype)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
@@ -74,14 +75,15 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
         beta = group['beta']
         store_triu_as_line = group['store_triu_as_line']
         q_dtype = getattr(torch, group['q_dtype'])
+        storage_dtype = getattr(torch, group['storage_dtype'])
 
         vals = []
 
-        for p, g in split_p_and_g_in_group(group):
+        for p, g in split_p_and_g_in_group(group, should_promote=False):
             state = self.state_(p)
 
             if 'Q' not in state:
-                state["exp_avg"] = torch.zeros_like(g)
+                state["exp_avg"] = torch.zeros_like(g, dtype=storage_dtype)
                 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
@@ -105,7 +107,9 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
 
         group["step"] += 1
 
-        torch._foreach_lerp_(exp_avg_list, grad_list, (1 - beta) / (1 - beta ** group["step"]))
+        stochastic_lerp_(exp_avg_list, grad_list, 1 - beta_debias(beta, group['step']))
+
+        lr = -warmup(lr, group['step'], group['warmup_steps'])
 
         grad_list, Q_list, Q_cache_list, exp_avg_list = list(grad_list), list(Q_list), list(Q_cache_list), list(
             exp_avg_list)
@@ -127,8 +131,4 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                     else:
                         torch.mul(q_.conj(), q_, out=c_)
 
-            set_(g, new)
-        grad_list = self.clip_fn(grad_list)
-
-        lr = -warmup(lr, group['step'], group['warmup_steps'])
-        update_param_(p_list, grad_list, lr, weight_decay)
+            update_param_([p], self.clip_fn([new]), lr, weight_decay)

heavyball/cached_psgd_kron.py

@@ -9,7 +9,7 @@ from typing import Optional
 import torch
 
 from .utils import update_param_, warmup, init_Q_exprs, trust_region_clip_, PSGDBase, split_p_and_g_in_group, \
-    line_to_triu, triu_to_line, set_, einsum_base, promote
+    line_to_triu, triu_to_line, einsum_base, promote, stochastic_lerp_, beta_debias
 
 
 class ForeachCachedPSGDKron(PSGDBase):
@@ -39,7 +39,8 @@ class ForeachCachedPSGDKron(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,  #
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,
+                 storage_dtype: str = 'float32',  #
                  # expert parameters
                  precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
@@ -56,7 +57,8 @@ class ForeachCachedPSGDKron(PSGDBase):
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
                         momentum_into_precond_update=momentum_into_precond_update, precond_lr=precond_lr,
                         precond_init_scale=precond_init_scale, step=0, warmup_steps=warmup_steps, merge_dims=merge_dims,
-                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
+                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype,
+                        storage_dtype=storage_dtype)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
@@ -71,15 +73,16 @@ class ForeachCachedPSGDKron(PSGDBase):
         beta = group['beta']
         store_triu_as_line = group['store_triu_as_line']
         q_dtype = getattr(torch, group['q_dtype'])
+        storage_dtype = getattr(torch, group['storage_dtype'])
         should_update = self.should_update(group)
 
         vals = []
 
-        for p, g in split_p_and_g_in_group(group):
+        for p, g in split_p_and_g_in_group(group, should_promote=False):
             state = self.state_(p)
 
             if 'Q' not in state:
-                state["exp_avg"] = torch.zeros_like(g)
+                state["exp_avg"] = torch.zeros_like(g, dtype=storage_dtype)
                 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
@@ -103,7 +106,9 @@ class ForeachCachedPSGDKron(PSGDBase):
 
         group["step"] += 1
 
-        torch._foreach_lerp_(exp_avg_list, grad_list, (1 - beta) / (1 - beta ** group["step"]))
+        stochastic_lerp_(exp_avg_list, grad_list, 1 - beta_debias(beta, group['step']))
+
+        lr = -warmup(lr, group['step'], group['warmup_steps'])
 
         grad_list, Q_list, Q_cache_list, exp_avg_list = list(grad_list), list(Q_list), list(Q_cache_list), list(
             exp_avg_list)
@@ -123,9 +128,5 @@ class ForeachCachedPSGDKron(PSGDBase):
                     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)
+            g = torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea)
+            update_param_([p], self.clip_fn([g]), lr, weight_decay)

heavyball/delayed_psgd.py

@@ -5,10 +5,10 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 """
 
 import torch
-from heavyball.utils import copy_stochastic_list_
 
+from heavyball.utils import stochastic_lerp_, beta_debias
 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_, promote
+    split_p_and_g_in_group, triu_to_line, line_to_triu, promote
 
 
 class ForeachDelayedPSGD(PSGDBase):
@@ -38,8 +38,8 @@ 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, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,  #
+                 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', #
                  # expert parameters
                  precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
@@ -55,12 +55,10 @@ class ForeachDelayedPSGD(PSGDBase):
         defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
                         momentum_into_precond_update=momentum_into_precond_update, precond_lr=precond_lr,
-                        precond_init_scale=precond_init_scale,
-                        step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
-                        store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
+                        precond_init_scale=precond_init_scale, step=0, warmup_steps=warmup_steps, merge_dims=merge_dims,
+                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
-
     def _step(self, group):
         should_update = self.should_update(group)
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
@@ -74,14 +72,15 @@ class ForeachDelayedPSGD(PSGDBase):
         beta = group['beta']
         store_triu_as_line = group['store_triu_as_line']
         q_dtype = getattr(torch, group['q_dtype'])
+        storage_dtype = getattr(torch, group['storage_dtype'])
 
         vals = []
 
-        for p, g in split_p_and_g_in_group(group):
+        for p, g in split_p_and_g_in_group(group, should_promote=False):
             state = self.state_(p)
 
             if 'Q' not in state:
-                state["exp_avg"] = torch.zeros_like(g)
+                state["exp_avg"] = torch.zeros_like(g, dtype=storage_dtype)
                 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
@@ -96,7 +95,9 @@ class ForeachDelayedPSGD(PSGDBase):
 
         group["step"] += 1
 
-        torch._foreach_lerp_(exp_avg_list, grad_list, (1 - beta) / (1 - beta ** group["step"]))
+        stochastic_lerp_(exp_avg_list, grad_list, beta_debias(beta, group["step"]))
+
+        lr = -warmup(lr, group['step'], group['warmup_steps'])
 
         Q_list, exp_avg_list = list(Q_list), list(exp_avg_list)
         for i, (p, g) in enumerate(zip(p_list, grad_list)):
@@ -106,10 +107,6 @@ class ForeachDelayedPSGD(PSGDBase):
             new = psgd_precond_grad(q, self.state_(p)["exprs"], ea)
             if should_update:
                 q32 = [promote(q_) for q_ in q]
-                self.do_update(group,[p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig], store_triu_as_line)
-            set_(g, new)
-
-        grad_list = self.clip_fn(grad_list)
-
-        lr = -warmup(lr, group['step'], group['warmup_steps'])
-        update_param_(p_list, grad_list, lr, weight_decay)
+                self.do_update(group, [p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig],
+                               store_triu_as_line)
+            update_param_([p], self.clip_fn([new]), lr, weight_decay)

heavyball/foreach_soap.py

@@ -1,7 +1,7 @@
 import torch
 
 from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, set_, \
-    split_p_and_g_in_group, StatefulOptimizer
+    split_p_and_g_in_group, StatefulOptimizer, exp_avg_
 
 
 class ForeachSOAP(StatefulOptimizer):
@@ -26,8 +26,7 @@ class ForeachSOAP(StatefulOptimizer):
                  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):
+                 split: bool = False, foreach: bool = True):
         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,
@@ -65,14 +64,12 @@ class ForeachSOAP(StatefulOptimizer):
         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'])
+        step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
+        denom = exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step_tensor)
 
         for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
             state = self.state_(p)

heavyball/p_adam.py

@@ -39,7 +39,7 @@ class ForeachPaLMPAdam(PSGDBase):
                  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,  #
+                 stochastic_schedule: bool = True,  storage_dtype:str ='float32',#
                  # expert parameters
                  precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
@@ -57,7 +57,7 @@ class ForeachPaLMPAdam(PSGDBase):
                         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)
+                        store_triu_as_line=store_triu_as_line, q_dtype=q_dtype, storage_dtype=storage_dtype)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
@@ -71,15 +71,16 @@ class ForeachPaLMPAdam(PSGDBase):
         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 split_p_and_g_in_group(group):
+        for p, g in split_p_and_g_in_group(group, should_promote=False):
             state = self.state_(p)
 
             if 'Q' not in state:
-                state['exp_avg'] = torch.zeros_like(g)
-                state['exp_avg_sq'] = torch.zeros_like(g)
+                state['exp_avg'] = torch.zeros_like(g, dtype=storage_dtype)
+                state['exp_avg_sq'] = torch.zeros_like(g, dtype=storage_dtype)
                 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
@@ -103,6 +104,8 @@ class ForeachPaLMPAdam(PSGDBase):
 
         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):
             psgd_precond_grad(Q, self.state_(p)["exprs"], g, inplace=True)
             ea = psgd_precond_grad(Q, self.state_(p)["exprs"], ea)
@@ -112,8 +115,5 @@ class ForeachPaLMPAdam(PSGDBase):
             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)
 
-        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/palm_foreach_soap.py

@@ -1,7 +1,8 @@
 import torch
 
 from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, set_, \
-    split_p_and_g_in_group, StatefulOptimizer
+    split_p_and_g_in_group, StatefulOptimizer, exp_avg_
+
 
 
 class PaLMForeachSOAP(StatefulOptimizer):
@@ -32,8 +33,7 @@ class PaLMForeachSOAP(StatefulOptimizer):
                  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):
+                 beta2_scale: float = 0.8, split: bool = False, foreach: bool = True):
         if betas[0] is not None:
             beta = betas[0]
         defaults = {"lr": lr, "beta": beta, "shampoo_beta": shampoo_beta, "eps": eps, "weight_decay": weight_decay,
@@ -75,13 +75,13 @@ class PaLMForeachSOAP(StatefulOptimizer):
         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'])
+        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)
+        denom = exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step_tensor)
 
         for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
             state = self.state_(p)

heavyball/precond_schedule_foreach_soap.py

@@ -2,8 +2,8 @@ import random
 
 import torch
 
-from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, \
-    precond_schedule, set_, split_p_and_g_in_group, StatefulOptimizer
+from .utils import init_preconditioner, update_preconditioner, project, beta_debias, update_param_, \
+    precond_schedule, set_, split_p_and_g_in_group, StatefulOptimizer, exp_avg_
 
 
 class PrecondScheduleForeachSOAP(StatefulOptimizer):
@@ -27,8 +27,7 @@ class PrecondScheduleForeachSOAP(StatefulOptimizer):
                  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,
-                 precond_scheduler=(1 / 3, 9), split: bool = False,
-                 foreach: bool = True):
+                 precond_scheduler=(1 / 3, 9), split: bool = False, foreach: bool = True):
         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,
@@ -68,14 +67,12 @@ class PrecondScheduleForeachSOAP(StatefulOptimizer):
         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'])
+        step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
+        denom = exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step_tensor)
 
         update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
         for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
@@ -89,8 +86,7 @@ class PrecondScheduleForeachSOAP(StatefulOptimizer):
             # 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)
+            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)

heavyball/precond_schedule_palm_foreach_soap.py

@@ -2,7 +2,7 @@ import random
 
 import torch
 
-from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, \
+from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_, update_param_, \
     precond_schedule, set_, split_p_and_g_in_group, StatefulOptimizer
 
 
@@ -81,9 +81,9 @@ class PrecondSchedulePaLMForeachSOAP(StatefulOptimizer):
 
         # 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'])
+        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)
+        denom = exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step_tensor)
 
         update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
         for p, g, ea, d in zip(p_list, grad, exp_avg, denom):

heavyball/precond_schedule_sfpsoap.py

@@ -2,8 +2,19 @@ import random
 
 import torch
 
-from .utils import init_preconditioner, update_preconditioner, project, set_, adaptive_gradient_clipping_, \
-    exp_avg_sq_, beta_debias, schedule_free_, warmup, ScheduleFree, precond_schedule, split_p_and_g_in_group
+from .utils import init_preconditioner, update_preconditioner, project, set_, adaptive_gradient_clipping_, exp_avg_sq_, \
+    beta_debias, schedule_free_, warmup, ScheduleFree, precond_schedule, split_p_and_g_in_group, copy_stochastic_list_, \
+    promote
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, eps):
+    eas32, gp32 = [list(map(promote, x)) for x in (exp_avg_sq, grad_projected)]
+    denom = exp_avg_sq_(eas32, gp32, old_debiased2, eps)
+    torch._foreach_div_(gp32, denom)
+
+    copy_stochastic_list_(exp_avg_sq, eas32)
+    copy_stochastic_list_(grad_projected, gp32)
 
 
 class PrecondScheduleSFPaLMSOAP(ScheduleFree):
@@ -40,8 +51,8 @@ class PrecondScheduleSFPaLMSOAP(ScheduleFree):
                  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, r=0.0,
-                 weight_lr_power=2.0, gradient_clip_val: float = 0.1, precond_scheduler=(1 / 3, 9),
-                 betas=(None, None), split: bool = False, foreach: bool = True):
+                 weight_lr_power=2.0, gradient_clip_val: float = 0.1, precond_scheduler=(1 / 3, 9), betas=(None, None),
+                 split: bool = False, foreach: bool = True):
         if betas[0] is not None:
             beta = betas[0]
         defaults = {"lr": lr, "beta": beta, "beta2_scale": beta2_scale, "eps": eps, "weight_decay": weight_decay,
@@ -103,8 +114,8 @@ class PrecondScheduleSFPaLMSOAP(ScheduleFree):
 
         # Decay the first and second moment running average coefficient
         # In-place operations to update the averages at the same time
-        denom = exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, group['eps'])
-        torch._foreach_div_(grad_projected, denom)
+        old_debiased_tensor = torch.empty((), dtype=torch.float32, device=p_list[0].device).fill_(old_debiased2)
+        _compilable_exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased_tensor, group["eps"])
 
         update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
 
@@ -114,13 +125,12 @@ class PrecondScheduleSFPaLMSOAP(ScheduleFree):
             # to the original space
             set_(gp, project(gp, state['Q'], back=True))
 
-            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2,
-                                  update_precond)
+            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2, update_precond)
 
         # Weight decay calculated at y
         if group["weight_decay"] > 0:
             torch._foreach_add_(grad, p_list, alpha=group["weight_decay"])
 
         lr = warmup(group['lr'], step, group['warmup_steps'])
-        group['weight_sum'] = schedule_free_(lr, group['weight_lr_power'], group['weight_sum'], group['beta'],
-                                             p_list, z, grad_projected, group['r'], step)
+        group['weight_sum'] = schedule_free_(lr, group['weight_lr_power'], group['weight_sum'], group['beta'], p_list,
+                                             z, grad_projected, group['r'], step)

heavyball/psgd_kron.py

@@ -9,7 +9,7 @@ from typing import Optional
 import torch
 
 from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust_region_clip_, PSGDBase, \
-    split_p_and_g_in_group, line_to_triu, triu_to_line, set_, promote
+    split_p_and_g_in_group, line_to_triu, triu_to_line, promote, stochastic_lerp_, beta_debias
 
 
 class ForeachPSGDKron(PSGDBase):
@@ -39,7 +39,8 @@ class ForeachPSGDKron(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,  #
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,
+                 storage_dtype: str = 'float32',  #
                  # expert parameters
                  precond_init_scale=1.0, precond_lr=0.1):
         if not 0.0 <= lr:
@@ -56,7 +57,7 @@ class ForeachPSGDKron(PSGDBase):
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
                         momentum_into_precond_update=momentum_into_precond_update, precond_lr=precond_lr,
                         precond_init_scale=precond_init_scale, step=0, warmup_steps=warmup_steps, merge_dims=merge_dims,
-                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
+                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype, storage_dtype=storage_dtype)
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
@@ -72,14 +73,15 @@ class ForeachPSGDKron(PSGDBase):
         beta = group['beta']
         store_triu_as_line = group['store_triu_as_line']
         q_dtype = getattr(torch, group['q_dtype'])
+        storage_dtype = getattr(torch, group['storage_dtype'])
 
         vals = []
 
-        for p, g in split_p_and_g_in_group(group):
+        for p, g in split_p_and_g_in_group(group, should_promote=False):
             state = self.state_(p)
 
             if 'Q' not in state:
-                state["exp_avg"] = torch.zeros_like(g)
+                state["exp_avg"] = torch.zeros_like(g, dtype=storage_dtype)
                 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
@@ -94,9 +96,14 @@ class ForeachPSGDKron(PSGDBase):
 
         group["step"] += 1
 
-        torch._foreach_lerp_(exp_avg_list, grad_list, (1 - beta) / (1 - beta ** group["step"]))
+        beta = beta_debias(beta, group["step"])
+        beta = torch.empty((), dtype=torch.float32, device=grad_list[0].device).fill_(1 - beta)
+        stochastic_lerp_(exp_avg_list, grad_list, 1 - beta)
 
         grad_list, Q_list, exp_avg_list = list(grad_list), list(Q_list), list(exp_avg_list)
+
+        lr = -warmup(lr, group['step'], group['warmup_steps'])
+
         for i, (p, g) in enumerate(zip(p_list, grad_list)):
             q_orig = Q_list.pop(0)
             ea = exp_avg_list.pop(0)
@@ -106,9 +113,5 @@ class ForeachPSGDKron(PSGDBase):
                 q32 = [promote(q_) for q_ in q]
                 self.do_update(group, [p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig],
                                store_triu_as_line)
-            set_(g, psgd_precond_grad(q, self.state_(p)["exprs"], 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)
+            g = psgd_precond_grad(q, self.state_(p)["exprs"], ea)
+            update_param_([p], self.clip_fn([g]), lr, weight_decay)

heavyball/pure_psgd.py

@@ -70,7 +70,7 @@ class ForeachPurePSGD(PSGDBase):
 
         vals = []
 
-        for p, g in split_p_and_g_in_group(group):
+        for p, g in split_p_and_g_in_group(group, should_promote=False):
             state = self.state_(p)
 
             if 'Q' not in state:
@@ -89,6 +89,7 @@ class ForeachPurePSGD(PSGDBase):
         group["step"] += 1
 
         Q_list = list(Q_list)
+        lr = -warmup(lr, group['step'], group['warmup_steps'])
         for i, (p, g) in enumerate(zip(p_list, grad_list)):
             q_orig = Q_list.pop(0)
             q = line_to_triu(q_orig) if store_triu_as_line else q_orig
@@ -97,8 +98,4 @@ class ForeachPurePSGD(PSGDBase):
                 q32 = [promote(q_) for q_ in q]
                 self.do_update(group, [p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
             psgd_precond_grad(q, self.state_(p)["exprs"], g, inplace=True)
-
-        grad_list = self.clip_fn(grad_list)
-
-        lr = -warmup(lr, group['step'], group['warmup_steps'])
-        update_param_(p_list, grad_list, lr, weight_decay)
+            update_param_([p], self.clip_fn([g]), lr, weight_decay)

heavyball/schedule_free_palm_foreach_soap.py

@@ -2,8 +2,18 @@ import random
 
 import torch
 
-from .utils import init_preconditioner, update_preconditioner, project, set_, adaptive_gradient_clipping_, \
-    exp_avg_sq_, beta_debias, schedule_free_, warmup, ScheduleFree, split_p_and_g_in_group
+from .utils import init_preconditioner, update_preconditioner, project, set_, adaptive_gradient_clipping_, exp_avg_sq_, \
+    beta_debias, schedule_free_, warmup, ScheduleFree, split_p_and_g_in_group, copy_stochastic_list_, promote
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, eps):
+    eas32, gp32 = [list(map(promote, x)) for x in (exp_avg_sq, grad_projected)]
+    denom = exp_avg_sq_(eas32, gp32, old_debiased2, eps)
+    torch._foreach_div_(gp32, denom)
+
+    copy_stochastic_list_(exp_avg_sq, eas32)
+    copy_stochastic_list_(grad_projected, gp32)
 
 
 class SFPaLMForeachSOAP(ScheduleFree):
@@ -95,8 +105,8 @@ class SFPaLMForeachSOAP(ScheduleFree):
 
         # Decay the first and second moment running average coefficient
         # In-place operations to update the averages at the same time
-        denom = exp_avg_sq_(exp_avg_sq, grad, new_debiased2, group["eps"])
-        torch._foreach_div_(grad_projected, denom)
+        old_debiased_tensor = torch.empty((), dtype=torch.float32, device=p_list[0].device).fill_(new_debiased2)
+        _compilable_exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased_tensor, group["eps"])
 
         update_precond = group['step'] > 0 and group['step'] % group['precondition_frequency'] == 0
 
@@ -107,13 +117,12 @@ class SFPaLMForeachSOAP(ScheduleFree):
             # CANT DO /= HERE AS EXP_AVG MAY POINT TO THE BUFFER
             set_(gp, project(gp, state['Q'], back=True))
 
-            update_preconditioner(g, state, max_precond_dim, precondition_1d, 1 - new_debiased2,
-                                  update_precond)
+            update_preconditioner(g, state, max_precond_dim, precondition_1d, 1 - new_debiased2, update_precond)
 
         # Weight decay calculated at y
         if group["weight_decay"] > 0:
             torch._foreach_add_(grad, p_list, alpha=group["weight_decay"])
 
         lr = warmup(group['lr'], step, group['warmup_steps'])
-        group['weight_sum'] = schedule_free_(lr, group['weight_lr_power'], group['weight_sum'], group['beta'],
-                                             p_list, z, grad_projected, group['r'], step)
+        group['weight_sum'] = schedule_free_(lr, group['weight_lr_power'], group['weight_sum'], group['beta'], p_list,
+                                             z, grad_projected, group['r'], step)

heavyball/utils.py

@@ -3,7 +3,7 @@ import gc
 import math
 import random
 import string
-from typing import List, Optional, Tuple, Callable
+from typing import List, Optional, Tuple, Callable, Union
 
 import numpy as np
 import torch
@@ -141,6 +141,7 @@ def beta_debias(beta, step):
     return 1 - (1 - beta) / (1 - beta ** step)
 
 
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
 def exp_avg_sq_(state, grad, beta2, eps, out=None):
     if isinstance(state, torch.Tensor):
         state.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
@@ -327,6 +328,26 @@ def get_orthogonal_matrix(mat):
     return final
 
 
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def stochastic_lerp_(x: List[torch.Tensor], y: List[torch.Tensor], a: Union[float, int, torch.Tensor]):
+    x32 = [promote(x_) for x_ in x]
+    y32 = [promote(y_) for y_ in y]
+
+    torch._foreach_lerp_(x32, y32, a)
+
+    copy_stochastic_list_(x, x32)
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def stochastic_add_(x: List[torch.Tensor], y: List[torch.Tensor], alpha: Union[float, int, torch.Tensor]):
+    x32 = [promote(x_) for x_ in x]
+    y32 = [promote(y_) for y_ in y]
+
+    [x_.add_(y_, alpha=alpha) for x_, y_ in zip(x32, y32)]
+
+    copy_stochastic_list_(x, x32)
+
+
 @decorator
 def compute_ggt(grad, GG, max_precond_dim, precondition_1d, beta):
     if grad.dim() == 1 and (not precondition_1d or grad.shape[0] > max_precond_dim):
@@ -409,9 +430,12 @@ def project(grad, Q, back: bool):
 
 
 class StatefulOptimizer(torch.optim.Optimizer):
-    def __init__(self, params, defaults, foreach: bool = True):
+    ema_decay: float = 0.001
+
+    def __init__(self, params, defaults, foreach: bool = True, use_ema: bool = False):
         super().__init__(params, {**defaults, 'foreach': foreach})
         self.fake_groups = {}
+        self.use_ema = use_ema
 
     def key(self, param: torch.Tensor):
         return (param.data_ptr(), tuple(param.shape))
@@ -445,6 +469,54 @@ class StatefulOptimizer(torch.optim.Optimizer):
     def _step(self, group):
         raise NotImplementedError
 
+    def ema_update(self):
+        with torch.no_grad():
+            for top_group in self.param_groups:
+                for group in self.get_groups(top_group):
+                    active_p = [p for p in group['params']]
+
+                    if not active_p:
+                        return
+
+                    k = group['ema_step'] = group.get('ema_step', -1) + 1
+
+                    for p in active_p:
+                        if 'param_ema' not in self.state_(p):
+                            self.state_(p)['param_ema'] = torch.zeros_like(p.data, memory_format=torch.preserve_format)
+
+                    y, param_ema = zip(*[(p.data, self.state_(p)['param_ema']) for p in active_p])
+                    torch._foreach_lerp_(param_ema, y, weight=beta_debias(1 - self.ema_decay, k + 1))
+
+    def copy_emas_to_params(self):
+        with torch.no_grad():
+            for top_group in self.param_groups:
+                for group in self.get_groups(top_group):
+                    active_p = [p for p in group['params']]
+
+                    if not active_p:
+                        return
+
+                    for p in active_p:
+                        if 'param_ema' in self.state_(p):
+                            p_clone = p.data.clone()
+                            set_(p.data, self.state_(p)['param_ema'])
+                            set_(self.state_(p)['param_ema'], p_clone)
+
+    def copy_params_to_emas(self):
+        with torch.no_grad():
+            for top_group in self.param_groups:
+                for group in self.get_groups(top_group):
+                    active_p = [p for p in group['params']]
+
+                    if not active_p:
+                        return
+
+                    for p in active_p:
+                        if 'param_ema' in self.state_(p):
+                            ema_clone = self.state_(p)['param_ema'].data.clone()
+                            set_(self.state_(p)['param_ema'], p.data)
+                            set_(p.data, ema_clone)
+
     def step(self, closure: Optional[Callable] = None):
         if closure is None:
             loss = None
@@ -455,6 +527,8 @@ class StatefulOptimizer(torch.optim.Optimizer):
             for top_group in self.param_groups:
                 for group in self.get_groups(top_group):
                     self._step(group)
+                    if self.use_ema:
+                        self.ema_update(group)
         return loss
 
 
@@ -497,6 +571,20 @@ def copy_stochastic_list_(target: List[torch.Tensor], source: List[torch.Tensor]
         copy_stochastic_(t, s)
 
 
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step):
+    beta1 = beta_debias(beta1, step)
+    beta2 = beta_debias(beta2, step)
+
+    g32, gp32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, grad_projected, exp_avg_sq]]
+
+    stochastic_lerp_(exp_avg, g32, 1 - beta1)
+    denom = exp_avg_sq_(exp_avg_sq32, gp32, beta2, 1e-8)
+
+    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
+    return denom
+
+
 # this can be dynamic for most optimizers - just not for PSGD. So, it's disabled for all
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True)
 def _compilable_copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
@@ -523,23 +611,26 @@ def copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
-def _compilable_update_one_(p, u, decay, add_fn, lr):
-    p32 = promote(p)
-    u32 = promote(u.view(p.shape))
+def _compilable_update_(p, u, decay, add_fn, lr):
+    u = [u_.view_as(p_) for u_, p_ in zip(u, p)]
+    p32, u32 = [list(map(promote, x)) for x in [p, u]]
+
     if decay > 0:
-        p32.mul_(1 - decay * lr)
-    if add_fn is None:
-        p32.add_(u32, alpha=lr)
-    else:
-        add_fn(p32, u32, lr)
-    copy_stochastic_(p, p32)
+        torch._foreach_mul_(p32, 1 - decay * lr)
+
+    for p32_, u32_ in zip(p32, u32):  # lr is data-dependent -> can't compile a foreach
+        if add_fn is None:
+            p32_.add_(u32_, alpha=lr)
+        else:
+            add_fn(p32_, u32_, lr)
+
+    copy_stochastic_list_(p, p32)
 
 
 def update_param_(param: List[torch.Tensor], update: List[torch.Tensor], lr: float, decay: float,
                   add_fn: callable = None):
     lr_tensor = torch.empty((), dtype=torch.float32, device=param[0].device).fill_(lr)
-    for p, u in zip(param, update):
-        _compilable_update_one_(p, u, decay, add_fn, lr_tensor)
+    _compilable_update_(param, update, decay, add_fn, lr_tensor)
 
 
 def precond_schedule(step, precond_scheduler, rng):
@@ -630,7 +721,7 @@ def init_Q_exprs(t, scale, max_size, min_ndim_triangular, memory_save_mode, dtyp
     return [Q, (exprA, tuple(exprGs), exprP)]
 
 
-@decorator
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
 def psgd_balance_Q(Q_in):
     norms = torch.stack([q.norm(float("inf")) for q in Q_in])
     geometric_mean = norms.log().mean().exp()
@@ -638,12 +729,14 @@ def psgd_balance_Q(Q_in):
     torch._foreach_mul_(Q_in, list(norms))
 
 
-def psgd_calc_A_and_conjB(exprA, G, Q, V):
-    md = min_dtype(Q)
-    A = torch.einsum(exprA, *[q.to(md) for q in Q], G.to(md))
+def psgd_calc_A_and_conjB(exprA, G, Q):
+    md = min_dtype(Q + [G])
+    A = torch.einsum(exprA, *[q.to(md) for q in Q], G.to(md)).to(G.dtype)
     order = G.dim()
     p = list(range(order))
-    conjB = torch.permute(V.conj(), p[1:] + p[:1])
+    V = torch.randn_like(G, dtype=promote(G.dtype))
+    conjB = torch.permute(V, p[1:] + p[:1])
+    Q = [promote(q) for q in Q]
     for i, q in enumerate(Q):
         if q.dim() <= 1:
             conjB /= q
@@ -651,7 +744,7 @@ def psgd_calc_A_and_conjB(exprA, G, Q, V):
             unsqueeze = conjB.dim() <= 1
             if unsqueeze:
                 conjB = conjB.unsqueeze(0)
-            conjB = torch.linalg.solve_triangular(q, conjB, upper=True, left=False, out=conjB)
+            conjB = torch.linalg.solve_triangular(q, conjB, upper=True, left=False)
             if unsqueeze:
                 conjB = conjB.squeeze(0)
         if i < order - 1:
@@ -661,33 +754,37 @@ def psgd_calc_A_and_conjB(exprA, G, Q, V):
 
 def psgd_lb(A, max_abs):
     A /= max_abs
-    aa = torch.real(A * A.conj())
-    value0, i = torch.max(torch.sum(aa, dim=0), 0)
-    value1, j = torch.max(torch.sum(aa, dim=1), 0)
+    a0 = torch.einsum('ij,ij->j', A, A)
+    a1 = torch.einsum('ij,ij->i', A, A)
+    value0 = torch.max(a0)
+    value1 = torch.max(a1)
+    i = torch.argmax(a0)
+    j = torch.argmax(a1)
 
-    ah = A.H
     comp = value0 > value1
-    x = torch.where(comp, A[:, i], A[j])
-    x = x.conj()
-    if x.dim() > 1:
-        x = torch.where(comp, x, x.T)
-    torch.matmul(x, torch.where(comp, A, A.T), out=x.view(1, -1))
-    x /= torch.linalg.vector_norm(x)
-    torch.matmul(x, torch.where(comp, ah, ah.T), out=x.view(1, -1))
-    x = torch.linalg.vector_norm(x)
+    x = torch.cond(comp, lambda a: torch.index_select(a, 1, i).flatten().contiguous(),  #
+                   lambda a: torch.index_select(a, 0, j).flatten().contiguous(), (A,))
+
+    x = torch.cond(comp, lambda x_, a: torch.einsum('i,ij->j', x_, a), lambda x_, a: torch.einsum('i,ji->j', x_, a),
+                   (x, A,))
+    x /= x.norm()
+    x = torch.cond(comp, lambda x_, a: torch.einsum('j,kj->k', x_, a), lambda x_, a: torch.einsum('j,jk->k', x_, a),
+                   (x, A,))
+    x = x.norm()
     x *= max_abs
     return x
 
 
-def psgd_update_precond(Q, exprs, V, G, step, tiny):
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+def psgd_update_precond(Q, exprs, G, precond_lr, tiny, oq, store_triu_as_line):
     """Update Kronecker product preconditioner Q with pair (V, G)."""
     exprA, exprGs, _ = exprs
 
-    A, conjB = psgd_calc_A_and_conjB(exprA, G, Q, V)
+    A, conjB = psgd_calc_A_and_conjB(exprA, G, Q)
 
-    for q, exprG in zip(Q, exprGs):
-        term1 = torch.einsum(exprG, A, A.conj())
-        term2 = torch.einsum(exprG, conjB.conj(), conjB)
+    for q, exprG, o in zip(Q, exprGs, oq):
+        term1 = promote(torch.einsum(exprG, A, A))
+        term2 = promote(torch.einsum(exprG, conjB, conjB))
 
         term2 += term1  # a + b
         term1 *= 2  # 2a
@@ -696,15 +793,19 @@ def psgd_update_precond(Q, exprs, V, G, step, tiny):
         else:
             term1 = term1 - term2
 
-        term1 *= step
+        term1 *= precond_lr
         norm = term2.norm(float('inf'))
         if q.dim() < 2:
-            term1 *= q
-            q.addcdiv_(term1, norm.clamp_(min=tiny), value=-1)
+            term1 *= q.to(term1.dtype)
+            term1 /= norm.clamp_(min=tiny)
         else:
             torch.triu(term1, out=term1)
-            term1 /= torch.where(norm > 0, psgd_lb(term2, norm), norm).clamp_(tiny)
-            q.addmm_(term1, q, alpha=-1)
+            term1 /= psgd_lb(term2, norm).clamp_(tiny)
+            torch.matmul(term1, q, out=term1)
+        if store_triu_as_line:
+            term1 = triu_to_line([term1])[0][1]
+            o = o[1]
+        stochastic_add_([o], [term1], -1)
 
 
 @decorator
@@ -838,18 +939,9 @@ class PSGDBase(StatefulOptimizer):
         group[name] = cumulative_prob + prob
         return int(group[name]) > int(cumulative_prob)
 
-    def do_update(self, group, p_list, grad_list, q_list, precond_lr, original_q: Optional[List] = None,
-                  store_triu_as_line=False):
-        if original_q:
-            if store_triu_as_line:
-                update_fn = update_triu_
-            else:
-                update_fn = copy_stochastic_list_
-        else:
-            update_fn = lambda x, y: None
-        for i, (p, grad, Q, oq) in enumerate(zip(p_list, grad_list, q_list, original_q)):
-            psgd_update_precond(Q, self.state_(p)["exprs"], torch.randn_like(grad), grad, precond_lr, self._tiny)
-            update_fn(oq, Q)
+    def do_update(self, group, p_list, grad_list, q_list, precond_lr, original_q: List, store_triu_as_line=False):
+        for p, grad, Q, oq in zip(p_list, grad_list, q_list, original_q):
+            psgd_update_precond(Q, self.state_(p)["exprs"], grad, precond_lr, self._tiny, oq, store_triu_as_line)
 
         if self.should_update(group, self.balance_probability, "balance_prob"):
             for g, q in zip(grad_list, original_q if original_q else q_list):
@@ -896,13 +988,19 @@ def merge_group(group, *tensors):
     return out
 
 
-def split_p_and_g_in_group(group: dict, skip_none: bool = True):
+def split_p_and_g_in_group(group: dict, skip_none: bool = True, should_promote: bool = True):
     for p in group["params"]:
         if skip_none and p.grad is None:
             continue
 
-        grad = None if p.grad is None else promote(p.grad)
-        p.grad = None
+        if p.grad is None:
+            grad = None
+        else:
+            if should_promote:
+                grad = promote(p.grad)
+            else:
+                grad = p.grad
+            p.grad = None
 
         p_views = merge_group(group, p)
         if grad is not None:

{heavyball-0.19.0.dist-info → heavyball-0.20.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.19.0
+Version: 0.20.0
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler
@@ -32,7 +32,7 @@ A simple package of efficient optimizers
 The goal is not to thrive for completeness, full maintenance or abstraction, but instead to provide a simple
 largely static alternative to `torch.optim` with more and better optimizers.
 
-Currently (2024-11-22, 0.19), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
+Currently (2024-11-22, 0.19.1), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
 recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psgd_efficiency.md)).
 
 ## Features

heavyball-0.20.0.dist-info/RECORD

@@ -0,0 +1,24 @@
+heavyball/__init__.py,sha256=iqP428JWwwx-XDOZ0nUdbCkOLEyfoqVyWZLQLAcwxaw,2214
+heavyball/cached_delayed_psgd_kron.py,sha256=apVzESMaQ8uxunHvfvYfyWA8HLbS25wQSd3j_YNEjGs,6603
+heavyball/cached_psgd_kron.py,sha256=3IETfsC0Ufu_8TPfo9SByGmztwjW6ktSFPwHNrUWkys,6601
+heavyball/delayed_psgd.py,sha256=0LaazbiBZOdx78EDS-945cW3bmeORjUvdFOGqdw3aMs,5631
+heavyball/foreach_adamw.py,sha256=Rb5U80cgUcEqlEbUU250UTWdoqA7nyiqkV5w1U4bWX4,2445
+heavyball/foreach_adopt.py,sha256=ecdi1fKg9i087OGjtKWVbE_DD6Yf4pvpzv4ELCcusvQ,3211
+heavyball/foreach_laprop.py,sha256=vi6C_gfjXxw5uN0KHgzxI9itUI1dcgOf3ufoO_VVMp0,2471
+heavyball/foreach_sfadamw.py,sha256=rLZORmCIMu9G09FdDgMSiI6pNq34IVoxsPVWtmeDdbQ,2753
+heavyball/foreach_soap.py,sha256=4mWSMWYTdjgiXiboI5DwdigecruDtNGKylGAFAVhCRA,4562
+heavyball/p_adam.py,sha256=J5QqFAlyLTQ1eQzM0LGxPdv4fEtZikIv9mJ_SSkO3ZY,6033
+heavyball/palm_foreach_sfadamw.py,sha256=JbNrcoquBGGUI5XNMFouDjpNurVHUW9DbX1A3tSrtno,3025
+heavyball/palm_foreach_soap.py,sha256=GzAwM8kOt1X0QCmUZDTdHwPxbJwjH8ic43dyAK5BYCA,6015
+heavyball/precond_schedule_foreach_soap.py,sha256=HcObXLfSNN_lKNb4nmC6tkdHcqDIMNX6hILpHKScqLc,4744
+heavyball/precond_schedule_palm_foreach_soap.py,sha256=xZ7CJvIfdu2RNAZt2g1S7Xb0Jyy1hNC4MowOFU3nWkk,6283
+heavyball/precond_schedule_sfpsoap.py,sha256=PNneiOkrRyV1yIZn91lPmYofd1_OiLqJTDy75RLpXJk,7273
+heavyball/psgd_kron.py,sha256=RSLJi5FnSmYjvYBufNDClnYRm-eN_Kpa1Ar2tNP6-X0,5824
+heavyball/pure_psgd.py,sha256=LZK0qmvZkBF8g00evaVLtW-sIUJmdoxag1K7O26AqEo,4820
+heavyball/schedule_free_palm_foreach_soap.py,sha256=_AqrnChY6iDlQUkF2YUxS7eLjSWCIuvEUOHvMHVM1yY,6873
+heavyball/utils.py,sha256=ESazD0yv14Aa8XKi_pz2CyfVkpcbgYcG2-WMvhQOnxk,35719
+heavyball-0.20.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.20.0.dist-info/METADATA,sha256=dJ43LOTrNqh7cDTDzZDSu57goP1gNhU3dfZ26BUK9hA,11926
+heavyball-0.20.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.20.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.20.0.dist-info/RECORD,,

heavyball-0.19.0.dist-info/RECORD

@@ -1,24 +0,0 @@
-heavyball/__init__.py,sha256=iqP428JWwwx-XDOZ0nUdbCkOLEyfoqVyWZLQLAcwxaw,2214
-heavyball/cached_delayed_psgd_kron.py,sha256=PQAER6UgVh5l87DGRZrJ8CVP9UhyCG5wJD9rPLnj_G8,6460
-heavyball/cached_psgd_kron.py,sha256=GaeneBp0irksCSBIrJY4D_0hCpZ-uSRPMhqVX_a-og8,6417
-heavyball/delayed_psgd.py,sha256=fhBWFLTSl1S2gHWCeYak-STaXRwpC56sWZGLFMKFEJM,5589
-heavyball/foreach_adamw.py,sha256=Rb5U80cgUcEqlEbUU250UTWdoqA7nyiqkV5w1U4bWX4,2445
-heavyball/foreach_adopt.py,sha256=ecdi1fKg9i087OGjtKWVbE_DD6Yf4pvpzv4ELCcusvQ,3211
-heavyball/foreach_laprop.py,sha256=vi6C_gfjXxw5uN0KHgzxI9itUI1dcgOf3ufoO_VVMp0,2471
-heavyball/foreach_sfadamw.py,sha256=rLZORmCIMu9G09FdDgMSiI6pNq34IVoxsPVWtmeDdbQ,2753
-heavyball/foreach_soap.py,sha256=h6ptMch7oaynvu3eIJtWnVXypDA_5JDVm3Zb3PNEma0,4634
-heavyball/p_adam.py,sha256=4zJDGJrpgUyVzr3GiELETFre4xr3-PE10OuAZj-jFM8,5883
-heavyball/palm_foreach_sfadamw.py,sha256=JbNrcoquBGGUI5XNMFouDjpNurVHUW9DbX1A3tSrtno,3025
-heavyball/palm_foreach_soap.py,sha256=g4hbiGRcti-J-a0SwAkP4ii5pU-aalsZH5bssyhroLk,5938
-heavyball/precond_schedule_foreach_soap.py,sha256=WLg5SzpJnKPZUvFyIvdwSZa1Umt5cpr3Kow_42orM-E,4863
-heavyball/precond_schedule_palm_foreach_soap.py,sha256=ammQrvRZFF-wc-wEiPEoFhS_7b8pdV61QfcLoQfimSo,6211
-heavyball/precond_schedule_sfpsoap.py,sha256=vq7jd302refKPa_9X2lkOTOtCCcTBVByPdojklrY8pA,6770
-heavyball/psgd_kron.py,sha256=wKjtI56iUnL5D8DseW60kxiXTAlMYNEf52CrvQaQMnI,5547
-heavyball/pure_psgd.py,sha256=iUy7mMKWxwNiVUMYrQ7SBnreu3t_XSbnhTW3a1yw4m0,4835
-heavyball/schedule_free_palm_foreach_soap.py,sha256=zkcikH5wWbzq4kOrmBjilvY3iWzuUddcv2HNEPKr3MI,6366
-heavyball/utils.py,sha256=BWscCHlGOw1_zfKYxNAAmfFeOXVpSJHuvqqlfL5A7_0,31690
-heavyball-0.19.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.19.0.dist-info/METADATA,sha256=1wORoS9rrjlug9tuJqXsbtVA9PphOBGcifiLRxmZNjs,11924
-heavyball-0.19.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.19.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.19.0.dist-info/RECORD,,