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

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.18.8.dist-info → heavyball-0.20.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.18.8
+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-21, 0.18.6), 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
@@ -45,8 +45,10 @@ recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psg
 * **ScheduleFree**: No learning rate schedule, but better convergence
 * [**Preconditioner Schedule**](https://github.com/lixilinx/psgd_torch/): Improved loss-per-step in early convergence,
   better step-per-second in late convergence (explained below)
-* **Memory-efficient storage** PSGD supports `store_triu_as_line` (default: `True`) to trade off memory usage for memory
-  bandwidth; turn it off for lower overheads (for more, see [PSGD Efficiency](docs/psgd_efficiency.md))
+* **Memory-efficient storage** PSGD supports `store_triu_as_line` (default: `True`) and `q_dtype` to trade off memory
+  usage for memory
+  bandwidth; Other optimizers have `storage_dtype`, supporting lower-precision EMAs at no(?) performance drop via
+  stochastic rounding
 
 ## Getting started
 
@@ -76,19 +78,19 @@ for _ in range(1000):
 
 ## Optimizers
 
-| Name                    | Description                                                                                                                                                       | Advantages / Disadvantages                                                                                                                                                                                                                                                                                                            |
-|-------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| **AdamW**               | More efficient (speed, memory) [AdamW](https://arxiv.org/abs/1711.05101)                                                                                          | + Faster than AdamW<br>+ Possibly more (numerically) stable                                                                                                                                                                                                                                                                           
-| **LaProp**              | More efficient (speed, memory) [LaProp](https://arxiv.org/abs/2002.04839)                                                                                         | + Same cost as AdamW<br>+ Marginally better converence (better proofs)<br>+ Higher hyperparameter stability<br>- Not a guaranteed win (can be neutral)<br>- No "Slingshot"                                                                                                                                                            |
-| **ADOPT**               | More efficient (speed, memory) [ADOPT](https://arxiv.org/abs/2411.02853)                                                                                          | + Same cost as AdamW<br>+ Rigorous mathematical convergence proofs, even for challenging models (GANs)<br>- Empirically underperforms LaProp<br>- no bf16                                                                                                                                                                             |
-| **SFAdamW**             | More efficient (speed, memory) [ScheduleFree AdamW](https://arxiv.org/abs/2405.15682)                                                                             | + Same cost as AdamW, but better eval perf<br>+ Full control over hyperparameters                                                                                                                                                                                                                                                     |
-| **PaLMSFAdamW**         | ForeachSFAdamW with [PaLM's beta2 schedule](https://arxiv.org/abs/2204.02311)                                                                                     | + Same cost as AdamW, but better eval perf<br>+ Less control, but faster early and more stable late convergence<br>+ ScheduleFree<br>- slow early convergence                                                                                                                                                                         |
-| **SOAP**                | More efficient (speed, memory) [SOAP](https://arxiv.org/abs/2409.11321)                                                                                           | + Faster convergence (loss-at-step)<br>+ Full control over hyperparameters<br>- more memory usage<br>- more hyperparameters<br>- higher overhead than AdamW (can be ammortized; better loss-at-second)                                                                                                                                |
-| **PaLMSOAP**            | ForeachSOAP with [PaLM's beta2 schedule](https://arxiv.org/abs/2204.02311)                                                                                        | + Faster convergence (loss-at-step)<br>+ Less control, but faster early and more stable late convergence<br>- more memory usage<br>- more hyperparameters<br>- higher overhead than AdamW (can be ammortized; better loss-at-second)                                                                                                  |
-| **SFPaLMSOAP**          | ScheduleFree PaLMForeachSOAP                                                                                                                                      | + Fast convergence (loss-at-step)<br>+ less memory usage than PaLMForeachSOAP (more tham AdamW)<br>- slower initial convergence than PaLMForeachSOAP (but allows higher LRs)<br>- higher overhead than AdamW (can be ammortized)                                                                                                      |
+| Name                          | Description                                                                                                                                                       | Advantages / Disadvantages                                                                                                                                                                                                                                                                                                            |
+|-------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| **AdamW**                     | More efficient (speed, memory) [AdamW](https://arxiv.org/abs/1711.05101)                                                                                          | + Faster than AdamW<br>+ Possibly more (numerically) stable                                                                                                                                                                                                                                                                           
+| **LaProp**                    | More efficient (speed, memory) [LaProp](https://arxiv.org/abs/2002.04839)                                                                                         | + Same cost as AdamW<br>+ Marginally better converence (better proofs)<br>+ Higher hyperparameter stability<br>- Not a guaranteed win (can be neutral)<br>- No "Slingshot"                                                                                                                                                            |
+| **ADOPT**                     | More efficient (speed, memory) [ADOPT](https://arxiv.org/abs/2411.02853)                                                                                          | + Same cost as AdamW<br>+ Rigorous mathematical convergence proofs, even for challenging models (GANs)<br>- Empirically underperforms LaProp<br>- no bf16                                                                                                                                                                             |
+| **SFAdamW**                   | More efficient (speed, memory) [ScheduleFree AdamW](https://arxiv.org/abs/2405.15682)                                                                             | + Same cost as AdamW, but better eval perf<br>+ Full control over hyperparameters                                                                                                                                                                                                                                                     |
+| **PaLMSFAdamW**               | ForeachSFAdamW with [PaLM's beta2 schedule](https://arxiv.org/abs/2204.02311)                                                                                     | + Same cost as AdamW, but better eval perf<br>+ Less control, but faster early and more stable late convergence<br>+ ScheduleFree<br>- slow early convergence                                                                                                                                                                         |
+| **SOAP**                      | More efficient (speed, memory) [SOAP](https://arxiv.org/abs/2409.11321)                                                                                           | + Faster convergence (loss-at-step)<br>+ Full control over hyperparameters<br>- more memory usage<br>- more hyperparameters<br>- higher overhead than AdamW (can be ammortized; better loss-at-second)                                                                                                                                |
+| **PaLMSOAP**                  | ForeachSOAP with [PaLM's beta2 schedule](https://arxiv.org/abs/2204.02311)                                                                                        | + Faster convergence (loss-at-step)<br>+ Less control, but faster early and more stable late convergence<br>- more memory usage<br>- more hyperparameters<br>- higher overhead than AdamW (can be ammortized; better loss-at-second)                                                                                                  |
+| **SFPaLMSOAP**                | ScheduleFree PaLMForeachSOAP                                                                                                                                      | + Fast convergence (loss-at-step)<br>+ less memory usage than PaLMForeachSOAP (more tham AdamW)<br>- slower initial convergence than PaLMForeachSOAP (but allows higher LRs)<br>- higher overhead than AdamW (can be ammortized)                                                                                                      |
 | **PrecondScheduleSFPaLMSOAP** | SFPaLMForeachSOAP with [preconditioner schedule](https://github.com/lixilinx/psgd_torch/), matching the error of PrecondEvery=2 with the cost of PrecondEvery=512 | + Better initial convergence than SFPaLMForeachSOAP<br>+ Significantly faster (sec/it) later<br>+ less memory usage than PaLMForeachSOAP (more tham AdamW)<br>- slower initial convergence than PaLMForeachSOAP (but allows higher LRs)<br>- higher overhead than AdamW (can be ammortized), goes to 0 with increasing number of step |
-| **PrecondSchedulePaLMSOAP** | PrecondScheduleSFPaLMForeachSOAP without schedule-free                                                                                                            | + Best initial convergence<br>+ Significantly faster (sec/it) later<br>+ high stability<br>- more memory usage than PrecondScheduleSFPaLMForeachSOAP<br>- higher overhead than AdamW (can be ammortized), goes to 0 with increasing number of steps                                                                                   |
-| **PrecondScheduleSOAP** | PrecondScheduleSFPaLMForeachSOAP without PaLM's beta2 schedule                                                                                                    | + Better initial convergence<br>+ Significantly faster (sec/it) later<br>- more memory usage than PrecondScheduleSFPaLMForeachSOAP<br>- higher overhead than AdamW (can be ammortized), goes to 0 with increasing number of steps                                                                                                     |
+| **PrecondSchedulePaLMSOAP**   | PrecondScheduleSFPaLMForeachSOAP without schedule-free                                                                                                            | + Best initial convergence<br>+ Significantly faster (sec/it) later<br>+ high stability<br>- more memory usage than PrecondScheduleSFPaLMForeachSOAP<br>- higher overhead than AdamW (can be ammortized), goes to 0 with increasing number of steps                                                                                   |
+| **PrecondScheduleSOAP**       | PrecondScheduleSFPaLMForeachSOAP without PaLM's beta2 schedule                                                                                                    | + Better initial convergence<br>+ Significantly faster (sec/it) later<br>- more memory usage than PrecondScheduleSFPaLMForeachSOAP<br>- higher overhead than AdamW (can be ammortized), goes to 0 with increasing number of steps                                                                                                     |
 
 ## Precond Schedule