heavyball 0.19.0__tar.gz → 0.21.0__tar.gz

{heavyball-0.19.0 → heavyball-0.21.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.19.0
+Version: 0.21.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.21.0), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
 recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psgd_efficiency.md)).
 
 ## Features

{heavyball-0.19.0 → heavyball-0.21.0}/README.md

@@ -8,7 +8,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.21.0), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
 recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psgd_efficiency.md)).
 
 ## Features

{heavyball-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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-0.19.0 → heavyball-0.21.0}/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,36 @@ def get_orthogonal_matrix(mat):
     return final
 
 
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_stochastic_lerp_(x: List[torch.Tensor], y: List[torch.Tensor], a: Union[float, int, torch.Tensor]):
+    for x_, y_ in zip(x, y):
+        x32 = promote(x_)
+        y32 = promote(y_)
+        x32.lerp_(y32, a)
+        copy_stochastic_(x_, x32)
+
+
+def stochastic_lerp_(x: List[torch.Tensor], y: List[torch.Tensor], a: Union[float, int, torch.Tensor]):
+    if not isinstance(a, torch.Tensor):
+        a = torch.empty((), dtype=torch.float32, device=x[0].device).fill_(a)
+    _compilable_stochastic_lerp_(x, y, a)
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_stochastic_add_(x: List[torch.Tensor], y: List[torch.Tensor], alpha: Union[float, int, torch.Tensor]):
+    for x_, y_ in zip(x, y):
+        x32 = promote(x_)
+        y32 = promote(y_)
+        x32.add_(y32, alpha=alpha)
+        copy_stochastic_(x_, x32)
+
+
+def stochastic_add_(x: List[torch.Tensor], y: List[torch.Tensor], alpha: Union[float, int, torch.Tensor]):
+    if not isinstance(alpha, torch.Tensor):
+        alpha = torch.empty((), dtype=torch.float32, device=x[0].device).fill_(alpha)
+    _compilable_stochastic_add_(x, y, alpha)
+
+
 @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 +440,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 +479,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 +537,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 +581,32 @@ 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 _compilable_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
+
+
+def exp_avg_(exp_avg: List[torch.Tensor], exp_avg_sq: List[torch.Tensor], grad: List[torch.Tensor],
+             grad_projected: List[torch.Tensor], beta1: float, beta2: float, step: int):
+    if isinstance(beta1, float):
+        beta1 = torch.empty((), dtype=torch.float32, device=exp_avg[0].device).fill_(beta1)
+    if isinstance(beta2, float):
+        beta2 = torch.empty((), dtype=torch.float32, device=exp_avg[0].device).fill_(beta2)
+    if isinstance(step, int):
+        step = torch.empty((), dtype=torch.int32, device=exp_avg[0].device).fill_(step)
+    denom = _compilable_exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step)
+    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 +633,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):
@@ -638,12 +751,13 @@ 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])
+    conjB = torch.randn(G.shape[1:] + G.shape[:1], dtype=promote(G.dtype), device=G.device)
+    Q = [promote(q) for q in Q]
     for i, q in enumerate(Q):
         if q.dim() <= 1:
             conjB /= q
@@ -651,7 +765,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 +775,29 @@ 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)
-
-    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)
+    a0 = torch.einsum('ij,ij->j', A, A)
+    i = torch.argmax(a0)
+
+    x = torch.index_select(A, 1, i).flatten().contiguous()
+
+    x = torch.einsum('i,ij->j', x, A)
+    x /= x.norm()
+    x = torch.einsum('j,kj->k', 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=True)
+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,18 +806,22 @@ 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
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
 def psgd_precond_grad(Q, exprs, G, inplace: bool = False):
     """Precondition gradient G with preconditioner Q."""
     md = min_dtype(Q)
@@ -838,18 +952,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 +1001,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 → heavyball-0.21.0}/heavyball.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.19.0
+Version: 0.21.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.21.0), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
 recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psgd_efficiency.md)).
 
 ## Features

{heavyball-0.19.0 → heavyball-0.21.0}/heavyball.egg-info/SOURCES.txt

@@ -29,6 +29,7 @@ test/test_bf16_params.py
 test/test_bf16_q.py
 test/test_bf16_storage.py
 test/test_closure.py
+test/test_ema.py
 test/test_foreach.py
 test/test_memory.py
 test/test_merge.py

{heavyball-0.19.0 → heavyball-0.21.0}/setup.py

@@ -10,7 +10,7 @@ setuptools.setup(
     name='heavyball',
     license='BSD',
     description='Efficient optimizers',
-    version='0.19.0',
+    version='0.21.0',
     long_description=README,
     url='https://github.com/clashluke/heavyball',
     packages=setuptools.find_packages(),

{heavyball-0.19.0 → heavyball-0.21.0}/test/test_bf16_params.py

@@ -20,10 +20,11 @@ def get_memory():
 
 @pytest.mark.parametrize("opt", heavyball.__all__)
 @pytest.mark.parametrize("size,depth", [(256, 2)])
-def test_foreach(opt, size, depth: int, iterations: int = 128, outer_iterations: int = 3):
+def test_foreach(opt, size, depth: int, iterations: int = 512, outer_iterations: int = 3):
     set_torch()
     opt = getattr(heavyball, opt)
 
+
     peaks = []
     losses = []
 

heavyball-0.21.0/test/test_ema.py

@@ -0,0 +1,61 @@
+import pytest
+import torch
+from torch import nn
+from torch._dynamo import config
+
+import heavyball
+import heavyball.utils
+from benchmark.utils import get_optim
+from heavyball.utils import clean, set_torch
+
+config.cache_size_limit = 128
+
+
+def get_memory():
+    clean()
+    torch.cuda.synchronize()
+    clean()
+    torch.cuda.synchronize()
+    return torch.cuda.memory_allocated()
+
+
+@pytest.mark.parametrize("opt", heavyball.__all__)
+@pytest.mark.parametrize("size,depth", [(256, 2)])
+def test_foreach(opt, size, depth: int, iterations: int = 128, outer_iterations: int = 3):
+    set_torch()
+    opt = getattr(heavyball, opt)
+
+    peaks = []
+    losses = []
+
+    for do_ema in [True, False]:
+        torch.manual_seed(0x2131290)
+        peaks.append([])
+        losses.append([])
+
+        for i in range(outer_iterations):
+            model = nn.Sequential(*[nn.Linear(size, size) for _ in range(depth)]).cuda()
+            o = get_optim(opt, model.parameters(), lr=1e-3)
+
+            for _ in range(iterations):
+                loss = model(torch.randn((1024, size), device='cuda')).square().mean()
+                loss.backward()
+                o.step()
+                o.zero_grad()
+                if do_ema:
+                    o.ema_update()
+                    o.copy_emas_to_params()
+                    o.copy_params_to_emas()
+                losses[-1].append(loss.detach())
+
+            if do_ema:
+                o.copy_emas_to_params()
+            loss = model(torch.randn((1024, size), device='cuda')).square().mean()
+            losses[-1].append(loss.detach())
+
+            del model, o
+            clean()
+
+    for i, (l0, l1) in enumerate(zip(*losses)):
+        print(i, l0.item(), l1.item())
+        assert l0.float() <= l1.float()