heavyball 0.25.1__py3-none-any.whl → 1.1.0__py3-none-any.whl

heavyball/utils.py

@@ -1,3 +1,11 @@
+"""
+
+
+Originally from Evan Walters and Omead Pooladzandi, 2024
+Modified under Creative Commons Attribution 4.0 International
+Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5ee8a1a4c29e4780bbf6c521e545189eff9/kron_torch/kron.py
+"""
+
 import functools
 import gc
 import math
@@ -15,7 +23,8 @@ from torch.utils._pytree import tree_map
 compile_mode = "max-autotune-no-cudagraphs"
 dynamic = False
 compile_mode_recommended_to_none = None
-zeroth_power_mode = 'qr'  # 'qr' is baseline, 'newtonschulz' converges better and faster, 'eigh' is perfect but slow
+zeroth_power_mode = 'qr'  # 'qr' is baseline, 'newtonschulz' converges better and faster
+tiny_bf16 = torch.finfo(torch.bfloat16).tiny
 
 
 def decorator(func):
@@ -60,30 +69,22 @@ def warmup(lr: float, step: int, warmup_steps: int):
 
 @decorator_knowngood
 def _compilable_schedule_free_(p: List[Tensor], z: List[Tensor], ckp1: Tensor, grad: List[Tensor], lr: Tensor,
-                               beta1: Tensor):
+                               beta1: Tensor, decay: float):
+    grad = [u_.view_as(p_) for u_, p_ in zip(grad, p)]
     p32, z32, g32 = [list(map(promote, x)) for x in (p, z, grad)]
     for p_, z_, g_ in zip(p32, z32, g32):
+        if decay != 0:
+            g_.add_(p_, alpha=decay)
         p_.lerp_(z_, ckp1)
-        p_.add_(g_, alpha=lr * (beta1 * (1 - ckp1) - 1))
-        z_.add_(g_, alpha=-lr)
+        p_.add_(g_, alpha=lr - lr * (beta1 * (1 - ckp1)))
+        z_.add_(g_, alpha=lr)
     copy_stochastic_list_(p, p32)
     copy_stochastic_list_(z, z32)
 
 
-def get_ckp1(lr, weight_lr_power, weight_sum, r, step):
-    weight = lr ** weight_lr_power * max(step, 1) ** r
-    weight_sum = weight_sum + weight
-
-    try:
-        ckp1 = weight / weight_sum
-    except ZeroDivisionError:
-        ckp1 = 0
-    return ckp1, weight_sum
-
-
 def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1: float, parameters: List[Tensor],
-                   z: List[Tensor], grad: list[Tensor], r: float = 0.0, step: int = 0):
-    weight = lr ** weight_lr_power * max(step, 1) ** r
+                   z: List[Tensor], grad: List[Tensor], r: float = 0.0, step: int = 0, decay: float = 0.0):
+    weight = abs(lr) ** weight_lr_power * max(step, 1) ** r
     weight_sum = weight_sum + weight
 
     try:
@@ -91,10 +92,9 @@ def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1:
     except ZeroDivisionError:
         ckp1 = 0
 
-    # These operations update y in-place,
-    # without computing x explicitly.
-    lr, ckp1 = scalar_guard(lr, parameters[0]), scalar_guard(ckp1, parameters[0])
-    _compilable_schedule_free_(parameters, z, ckp1, grad, lr, beta1)
+    grad, parameters, z = list_guard(grad, parameters, z)
+    lr, ckp1, beta1 = scalar_guard(lr, ckp1, beta1, grad[0])
+    _compilable_schedule_free_(parameters, z, ckp1, grad, lr, beta1, decay)
     return weight_sum
 
 
@@ -162,10 +162,13 @@ def beta_debias(beta, step):
 @decorator_knowngood
 def _compilable_exp_avg_sq_(state: List[Tensor], grad: List[Tensor], beta2: Tensor, eps: Tensor,
                             out: List[Optional[Tensor]]):
-    torch._foreach_mul_(state, beta2)
-    [s.addcmul_(g, g, value=1 - beta2) for s, g in zip(state, grad)]
-    denom = torch._foreach_sqrt(state)
-    [denom.clamp_(min=eps) for denom in denom]
+    s32, g32 = [list(map(promote, x)) for x in (state, grad)]
+    torch._foreach_mul_(s32, beta2)
+    [s.addcmul_(g, g, value=1 - beta2) for s, g in zip(s32, g32)]
+    denom = torch._foreach_sqrt(s32)
+    [d.clamp_(min=eps) for d in denom]
+    copy_stochastic_list_(state, s32)
+
     if out[0] is None:
         return denom
 
@@ -174,15 +177,32 @@ def _compilable_exp_avg_sq_(state: List[Tensor], grad: List[Tensor], beta2: Tens
 
 
 def exp_avg_sq_(state, grad, beta2, eps, out=None):
-    state, grad, out = list_guard(state), list_guard(grad), list_guard(out)
-    beta2, eps = scalar_guard(beta2, state[0]), scalar_guard(eps, state[0])
+    state, grad, out = list_guard(state, grad, out)
+    beta2, eps = scalar_guard(beta2, eps, state[0])
     return _compilable_exp_avg_sq_(state, grad, beta2, eps, out)
 
 
-def adaptive_gradient_clipping_(parameters: List[Tensor], gradients: List[Tensor], clip_val: float,
-                                minimum: float = 1e-3, eps: float = 1e-8):
-    if clip_val <= 0:
-        return
+@decorator_knowngood
+def _compilable_scale_by_exp_avg_sq_(state: List[Tensor], grad: List[Tensor], beta2: Tensor, eps: Tensor):
+    s32, g32 = [list(map(promote, x)) for x in (state, grad)]
+    torch._foreach_mul_(s32, beta2)
+    [s.addcmul_(g, g, value=1 - beta2) for s, g in zip(s32, g32)]
+    denom = torch._foreach_sqrt(s32)
+    [d.clamp_(min=eps) for d in denom]
+    out = torch._foreach_div_(g32, denom)
+    copy_stochastic_list_(state, s32)
+    copy_stochastic_list_(grad, out)
+
+
+def scale_by_exp_avg_sq_(exp_avg_sq, grad, beta2, eps):
+    grad, exp_avg_sq = list_guard(grad, exp_avg_sq)
+    beta2, eps = scalar_guard(beta2, eps, grad[0])
+    _compilable_scale_by_exp_avg_sq_(exp_avg_sq, grad, beta2, eps)
+    return grad
+
+
+@decorator_knowngood
+def _compilable_agc_(parameters: List[Tensor], gradients: List[Tensor], clip_val: float, minimum: float, eps: float):
     p_norm = torch._foreach_norm(parameters)
     g_norm = torch._foreach_norm(gradients)
     torch._foreach_maximum_(p_norm, minimum)
@@ -190,7 +210,16 @@ def adaptive_gradient_clipping_(parameters: List[Tensor], gradients: List[Tensor
     torch._foreach_div_(p_norm, g_norm)
     torch._foreach_mul_(p_norm, clip_val)
     torch._foreach_minimum_(p_norm, 1)
-    torch._foreach_mul_(gradients, p_norm)
+    return torch._foreach_mul(gradients, p_norm)
+
+
+def adaptive_gradient_clipping_(parameters: List[Tensor], gradients: List[Tensor], clip_val: float,
+                                minimum: float = 1e-3, eps: float = 1e-8):
+    if clip_val <= 0:
+        return gradients
+    parameters, gradients = list_guard(parameters, gradients)
+    clip_val = scalar_guard(clip_val, parameters[0])
+    return _compilable_agc_(parameters, gradients, clip_val, minimum, eps)
 
 
 def is_compiling():
@@ -205,10 +234,7 @@ def set_(dst: Tensor, src: Tensor):
         return
     if src.shape != dst.shape:
         src = src.reshape_as(dst)
-    if not is_compiling() and src.is_contiguous() and dst.is_contiguous() and src.dtype == dst.dtype:
-        dst.set_(src)
-    else:
-        dst.copy_(src)
+    dst.copy_(src)
 
 
 def clean():
@@ -226,33 +252,29 @@ def set_torch():
 
 
 @decorator
-def zeropower_via_newtonschulz5(G, init, steps=2, eps=1e-7):
+def zeropower_via_newtonschulz5(G, steps=5, eps=1e-7):
     """
-    Modified from "modded-nanogpt" under the MIT license:
-    Original: https://github.com/KellerJordan/modded-nanogpt/blob/a0dcbfdd9a0617d091d5123cfc354745428e40d3/train_gpt2.py
-
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
     quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
     of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
     zero even beyond the point where the iteration no longer converges all the way to one everywhere
     on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
-    where S' is diagonal with S_{ii}' \sim Uniform(0.5, 1.5), which turns out not to hurt model
+    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
     a, b, c = (3.4445, -4.7750, 2.0315)
-    X = G.float()
-    init = init / (init.norm() + eps)  # ensure top singular value <= 1
-    X = X / (X.norm() + eps)  # ensure top singular value <= 1
+    X = G.bfloat16()
+    X /= (X.norm() + eps)  # ensure top singular value <= 1
     if G.size(0) > G.size(1):
         X = X.T
     for _ in range(steps):
-        A = X @ X.T  # preconditioner
-        B = A @ init
-        init = X = a * init + b * B + c * A @ B
+        A = X @ X.T
+        B = b * A + c * A @ A  # adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+        X = a * X + B @ X
     if G.size(0) > G.size(1):
         X = X.T
-    return X
+    return X.to(G.dtype)
 
 
 def ortho(x):
@@ -264,6 +286,53 @@ def ortho(x):
     raise NotImplementedError(f"Unknown zeroth_power_mode: {zeroth_power_mode}")
 
 
+@decorator_knowngood
+def _compilable_heavyball_momentum_(state, grad, beta):
+    s32, g32 = [list(map(promote, x)) for x in (state, grad)]
+    torch._foreach_mul_(s32, beta)
+    torch._foreach_add_(s32, g32)
+    copy_stochastic_list_(state, s32)
+    copy_stochastic_list_(grad, s32)
+
+
+@decorator_knowngood
+def _compilable_nesterov_momentum_(state, grad, beta):
+    s32, g32 = [list(map(promote, x)) for x in (state, grad)]
+    torch._foreach_mul_(s32, beta)
+    torch._foreach_add_(s32, g32)
+    [g.add_(s, alpha=beta) for g, s in zip(g32, s32)]
+    copy_stochastic_list_(state, s32)
+    copy_stochastic_list_(grad, g32)
+
+
+def heavyball_momentum(state, grad, beta):
+    state, grad = list_guard(state, grad)
+    beta = scalar_guard(beta, state[0])
+    _compilable_heavyball_momentum_(state, grad, beta)
+    return grad
+
+
+def nesterov_momentum(state, grad, beta):
+    state, grad = list_guard(state, grad)
+    beta = scalar_guard(beta, state[0])
+    _compilable_nesterov_momentum_(state, grad, beta)
+    return grad
+
+
+@decorator_knowngood
+def inplace_orthogonal_(x, mode, out):
+    if mode == 'qr':
+        y = torch.linalg.qr(x).Q
+    elif mode == 'svd':
+        u, s, v = torch.linalg.svd(x)
+        y = u @ v.T
+    elif mode == 'newtonschulz':
+        y = zeropower_via_newtonschulz5(x, 5)
+    else:
+        raise NotImplementedError(f"Unknown zeroth_power_mode: {mode}")
+    set_(out, y)
+
+
 def get_orthogonal_matrix_QR(GG, Q, exp_avg_sq):
     """
     Computes the eigenbases of the preconditioner using one round of power iteration
@@ -294,17 +363,7 @@ def get_orthogonal_matrix_QR(GG, Q, exp_avg_sq):
         est_eig = torch.einsum('ij,ij->j', o, tmp)
         sort_idx = torch.argsort(est_eig, descending=True)
         indices.append(sort_idx)
-        if zeroth_power_mode == 'eigh':
-            set_(q, torch.linalg.eigh(m)[1])
-        elif zeroth_power_mode.startswith('newtonschulz'):
-            iterations = zeroth_power_mode[len('newtonschulz'):]
-            if iterations == '':
-                iterations = 10
-            else:
-                iterations = int(iterations)
-            set_(q, zeropower_via_newtonschulz5(m, o[:, sort_idx], iterations))
-        else:
-            set_(q, ortho(tmp[:, sort_idx]))
+        inplace_orthogonal_(tmp[:, sort_idx], q)
 
     indices = tuple(slice(None) if ind is None else ind.view(*(1,) * i, -1, *(1,) * (exp_avg_sq.dim() - i - 1))  #
                     for i, ind in enumerate(indices))
@@ -353,8 +412,6 @@ def get_orthogonal_matrix(mat):
 
         Q = torch.flip(Q, [1])
 
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
         final.append(Q)
 
     return final
@@ -369,24 +426,57 @@ def _compilable_stochastic_lerp_(x: List[Tensor], y: List[Tensor], a: Union[floa
         copy_stochastic_(x_, x32)
 
 
+def get_beta1(group):
+    beta = None
+    if 'beta' in group:
+        beta = group['beta']
+    if beta is None and 'betas' in group:
+        beta = group['betas'][0]
+    if beta is None:
+        raise ValueError("Beta not found in group.")
+    return beta
+
+
+def get_beta2(group):
+    if 'beta2_scale' in group:
+        step = max(group.get("step", 1), 1)
+        return 1 - step ** -group['beta2_scale']
+    if 'betas' in group:
+        return group['betas'][1]
+    raise ValueError("Beta2 not found in group.")
+
+
 def stochastic_lerp_(x: List[Tensor], y: List[Tensor], a: Union[float, int, Tensor]):
-    x, y = list_guard(x), list_guard(y)
+    x, y = list_guard(x, y)
     a = scalar_guard(a, x[0])
     _compilable_stochastic_lerp_(x, y, a)
 
 
-def list_guard(x):
-    if isinstance(x, (list, tuple)):
-        return x
-    return [x]
+def list_guard(*xs):
+    out = []
+    for x in xs:
+        if isinstance(x, (list, tuple)):
+            out.append(x)
+        else:
+            out.append([x])
+    if len(xs) == 1:
+        return out[0]
+    return out
 
 
-def scalar_guard(x, ref):
-    if isinstance(x, float):
-        return torch.empty((), dtype=torch.float32, device=ref.device).fill_(x)
-    if isinstance(x, int):
-        return torch.empty((), dtype=torch.int64, device=ref.device).fill_(x)
-    return x
+def scalar_guard(*args):
+    *xs, ref = args
+    out = []
+    for x in xs:
+        if isinstance(x, float):
+            out.append(torch.empty((), dtype=torch.float32, device=ref.device).fill_(x))
+        elif isinstance(x, int):
+            out.append(torch.empty((), dtype=torch.int64, device=ref.device).fill_(x))
+        else:
+            out.append(x)
+    if len(xs) == 1:
+        return out[0]
+    return out
 
 
 @decorator_knowngood
@@ -399,7 +489,7 @@ def _compilable_stochastic_add_(x: List[Tensor], y: List[Tensor], alpha: Union[f
 
 
 def stochastic_add_(x: List[Tensor], y: List[Tensor], alpha: Union[float, int, Tensor]):
-    x, y = list_guard(x), list_guard(y)
+    x, y = list_guard(x, y)
     alpha = scalar_guard(alpha, x[0])
     _compilable_stochastic_add_(x, y, alpha)
 
@@ -435,35 +525,35 @@ def min_dtype(xs: List[Tensor]):
     return torch.float32
 
 
-def update_preconditioner(grad, state, max_precond_dim, precondition_1d, beta, update_precond):
+def update_preconditioner(grad, Q, GG, exp_avg_sq, max_precond_dim, precondition_1d, beta, update_precond):
     """
     Updates the preconditioner matrices and the eigenbases (L, R, Q_L, Q_R in the paper).
     """
-    compute_ggt(grad, state['GG'], max_precond_dim, precondition_1d, beta)
-    if state['Q'] is None:
-        state['Q'] = get_orthogonal_matrix(state['GG'])
+    compute_ggt(grad, GG, max_precond_dim, precondition_1d, beta)
     if update_precond:
-        get_orthogonal_matrix_QR(state['GG'], state['Q'], state['exp_avg_sq'])
+        get_orthogonal_matrix_QR(GG, Q, exp_avg_sq)
 
 
-def init_preconditioner(grad, state, max_precond_dim=10000, precondition_1d=False):
+def init_preconditioner(grad, state, beta, max_precond_dim=10000, precondition_1d=False):
     """
     Initializes the preconditioner matrices (L and R in the paper).
     """
-    state['Q'] = None  # Will hold all the eigenbases of the preconditioner.
     state['GG'] = []  # Will hold all the preconditioner matrices (L and R in the paper).
     if grad.dim() == 1:
-        if not precondition_1d or grad.shape[0] > max_precond_dim:
+        if precondition_1d or grad.shape[0] > max_precond_dim:
+            state['GG'].append(torch.zeros(grad.shape[0], grad.shape[0], device=grad.device, dtype=grad.dtype))
+        else:
             state['GG'].append([])
-            return
-        state['GG'].append(torch.zeros(grad.shape[0], grad.shape[0], device=grad.device, dtype=grad.dtype))
-        return
 
-    for sh in grad.shape:
-        if sh > max_precond_dim:
-            state['GG'].append([])
-        else:
-            state['GG'].append(torch.zeros(sh, sh, device=grad.device, dtype=grad.dtype))
+    else:
+        for sh in grad.shape:
+            if sh > max_precond_dim:
+                state['GG'].append([])
+            else:
+                state['GG'].append(torch.zeros(sh, sh, device=grad.device, dtype=grad.dtype))
+
+    compute_ggt(grad, state['GG'], max_precond_dim, precondition_1d, beta)
+    state['Q'] = get_orthogonal_matrix(state['GG'])
 
 
 @decorator
@@ -629,74 +719,87 @@ class StatefulOptimizer(torch.optim.Optimizer):
         return loss
 
 
-
-class ScheduleFree(StatefulOptimizer):
-    def eval(self):
-        for group in self.param_groups:
-            train_mode = group['train_mode']
-            beta1 = group['beta'] if 'beta' in group else group['betas'][0]
-            if beta1 > 0 and train_mode:
-                for p in group['params']:
-                    state = self.state_(p)
-                    if 'z' in state:
-                        # Set p.data to x
-                        z = promote(state['z'])
-                        p32 = promote(p.data)
-                        p32.lerp_(end=z, weight=1 - 1 / beta1)
-                        copy_stochastic_(p.data, p32)
-                group['train_mode'] = False
-
-    def train(self):
-        for group in self.param_groups:
-            train_mode = group['train_mode']
-            beta1 = group['beta'] if 'beta' in group else group['betas'][0]
-            if beta1 > 0 and not train_mode:
-                for p in group['params']:
-                    state = self.state_(p)
-                    if 'z' in state:
-                        z = promote(state['z'])
-                        p32 = promote(p.data)
-                        p32.lerp_(end=z, weight=1 - beta1)
-                        copy_stochastic_(p.data, p32)
-                group['train_mode'] = True
-
-    def _step(self):
-        raise NotImplementedError
-
-
 def copy_stochastic_list_(target: List[Tensor], source: List[Tensor]):
     for t, s in zip(target, source):
         copy_stochastic_(t, s)
 
 
 @decorator_knowngood
-def _compilable_exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor],
-                         grad_projected: List[Tensor], beta1: Tensor, beta2: Tensor, step: Tensor):
+def _compilable_adam_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: Tensor, beta2: Tensor,
+                      step: Tensor):
     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]]
+    g32, exp_avg_sq32, exp_avg32 = [list(map(promote, x)) for x in [grad, exp_avg_sq, exp_avg]]
 
-    stochastic_lerp_(exp_avg, g32, 1 - beta1)
-    denom = exp_avg_sq_(exp_avg_sq32, gp32, beta2, 1e-8)
+    [ea32.lerp_(g, 1 - beta1) for ea32, g in zip(exp_avg32, g32)]
+    denom = exp_avg_sq_(exp_avg_sq32, g32, beta2, 1e-8)
+    u32 = torch._foreach_div(exp_avg32, denom)
+
+    copy_stochastic_list_(exp_avg, exp_avg32)
+    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
+    copy_stochastic_list_(grad, u32)
+
+
+def adam_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: float, beta2: float, step: int):
+    exp_avg, exp_avg_sq, grad = map(list_guard, (exp_avg, exp_avg_sq, grad))
+    beta1, beta2, step = scalar_guard(beta1, beta2, step, exp_avg[0])
+    _compilable_adam_(exp_avg, exp_avg_sq, grad, beta1, beta2, step)
+    return grad
 
+
+@decorator_knowngood
+def _fused_compilable_adam_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor],
+                            beta1: Tensor, beta2: Tensor, step: Tensor, decay: Tensor, lr: Tensor, eps: Tensor,
+                            caution: bool):
+    beta1 = beta_debias(beta1, step)
+    beta2 = beta_debias(beta2, step)
+
+    g32, exp_avg_sq32, exp_avg32 = [list(map(promote, x)) for x in [grad, exp_avg_sq, exp_avg]]
+
+    [ea32.lerp_(g, 1 - beta1) for ea32, g in zip(exp_avg32, g32)]
+    denom = exp_avg_sq_(exp_avg_sq32, g32, beta2, 1e-8)
+    u32 = torch._foreach_div(exp_avg32, denom)
+
+    copy_stochastic_list_(exp_avg, exp_avg32)
     copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
-    return denom
+    _compilable_update_(y, u32, decay, lambda a, b, c: a.add_(b, alpha=c), lr, caution, g32)
+
+
+def fused_adam_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: float,
+                beta2: float, step: int, lr: float, eps: float, decay: float, caution: bool):
+    y, exp_avg, exp_avg_sq, grad = list_guard(y, exp_avg, exp_avg_sq, grad)
+    beta1, beta2, step, lr = scalar_guard(beta1, beta2, step, lr, y[0])
+    return _fused_compilable_adam_(y, exp_avg, exp_avg_sq, grad, beta1, beta2, step, decay, lr, eps, caution)
+
 
+@decorator_knowngood
+def _compilable_laprop_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: Tensor,
+                        beta2: Tensor, step: Tensor):
+    beta1 = beta_debias(beta1, step)
+    beta2 = beta_debias(beta2, step)
+
+    gp32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg_sq]]
+
+    denom = exp_avg_sq_(exp_avg_sq32, gp32, beta2, 1e-8)
+    gp32 = torch._foreach_div(gp32, denom)
+    stochastic_lerp_(exp_avg, gp32, 1 - beta1)
+
+    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
+    copy_stochastic_list_(grad, exp_avg)
 
-def exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], grad_projected: List[Tensor],
-             beta1: float, beta2: float, step: int):
-    exp_avg, exp_avg_sq, grad, grad_projected = list_guard(exp_avg), list_guard(exp_avg_sq), list_guard(
-        grad), list_guard(grad_projected)
-    beta1, beta, step = scalar_guard(beta1, exp_avg[0]), scalar_guard(beta2, exp_avg[0]), scalar_guard(step, exp_avg[0])
-    denom = _compilable_exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step)
-    return denom
 
+def laprop_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: float, beta2: float, step: int):
+    exp_avg, exp_avg_sq, grad = list_guard(exp_avg, exp_avg_sq, grad)
+    beta1, beta2, step = scalar_guard(beta1, beta2, step, exp_avg[0])
+    _compilable_laprop_(exp_avg, exp_avg_sq, grad, beta1, beta2, step)
+    return grad
 
 
 @decorator_knowngood
-def _compilable_laprop_exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor],
-                         grad_projected: List[Tensor], beta1: Tensor, beta2: Tensor, step: Tensor):
+def _fused_compilable_laprop_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor],
+                              grad_projected: List[Tensor], beta1: Tensor, beta2: Tensor, step: Tensor, lr: Tensor,
+                              decay: Tensor, caution: bool):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -705,31 +808,89 @@ def _compilable_laprop_exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor],
     denom = exp_avg_sq_(exp_avg_sq32, gp32, beta2, 1e-8)
     gp32 = torch._foreach_div(gp32, denom)
     stochastic_lerp_(exp_avg, gp32, 1 - beta1)
+    update_param_(y, gp32, lr, decay, caution=caution, grad=gp32)
 
     copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
 
 
-def laprop_exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad_projected: List[Tensor],
-             beta1: float, beta2: float, step: int):
-    exp_avg, exp_avg_sq, grad_projected = list_guard(exp_avg), list_guard(exp_avg_sq), list_guard(grad_projected)
-    beta1, beta, step = scalar_guard(beta1, exp_avg[0]), scalar_guard(beta2, exp_avg[0]), scalar_guard(step, exp_avg[0])
-    _compilable_laprop_exp_avg_(exp_avg, exp_avg_sq, grad_projected, beta1, beta2, step)
+def fused_laprop_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: float,
+                  beta2: float, step: int, lr: float, decay: float, caution: bool):
+    exp_avg, exp_avg_sq, grad, y = list_guard(exp_avg, exp_avg_sq, grad, y)
+    beta1, beta2, step, lr = scalar_guard(beta1, beta2, step, lr, exp_avg[0])
+    _fused_compilable_laprop_(y, exp_avg, exp_avg_sq, grad, beta1, beta2, step, lr, decay, caution)
 
 
 @decorator_knowngood
-def _compilable_copy_stochastic_(target: Tensor, source: Tensor):
-    """Taken as-is from https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905"""
-    # create a random 16 bit integer
-    result = torch.randint_like(source, dtype=torch.int32, low=0, high=(1 << 16))
+def _fused_compilable_adopt_(y, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay, caution):
+    g32, exp_avg32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg, exp_avg_sq]]
+    update_param_(y, exp_avg, lr, decay, caution=caution, grad=g32)
+
+    beta1 = beta_debias(beta1, step)
+    denom = torch._foreach_sqrt(exp_avg_sq32)
+    [denom.clamp_(min=eps) for denom in denom]
+    torch._foreach_mul_(exp_avg32, beta1)
+    [ea32.addcdiv_(g, d, value=1 - beta1) for ea32, g, d in zip(exp_avg32, g32, denom)]
+
+    beta2 = beta_debias(beta2, step + 1)
+    torch._foreach_mul_(exp_avg_sq32, beta2)
+    [eas32.addcmul_(g, g, value=1 - beta2) for eas32, g in zip(exp_avg_sq32, g32)]
+
+    copy_stochastic_list_(exp_avg, exp_avg32)
+    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
+
+
+def fused_adopt_(y, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay, caution):
+    exp_avg, exp_avg_sq, grad, y = list_guard(exp_avg, exp_avg_sq, grad, y)
+    beta1, beta2, step, lr = scalar_guard(beta1, beta2, step, lr, exp_avg[0])
+    _fused_compilable_adopt_(y, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay, caution)
+
+
+@decorator_knowngood
+def _compilable_adopt_(grad, exp_avg_sq, exp_avg, beta1, beta2, step):
+    g32, exp_avg32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg, exp_avg_sq]]
+    update = [e.clone() for e in exp_avg]
+
+    beta1 = beta_debias(beta1, step)
+    denom = torch._foreach_sqrt(exp_avg_sq32)
+    [denom.clamp_(min=1e-8) for denom in denom]
+    torch._foreach_mul_(exp_avg32, beta1)
+    [ea32.addcdiv_(g, d, value=1 - beta1) for ea32, g, d in zip(exp_avg32, g32, denom)]
+
+    beta2 = beta_debias(beta2, step + 1)
+    torch._foreach_mul_(exp_avg_sq32, beta2)
+    [eas32.addcmul_(g, g, value=1 - beta2) for eas32, g in zip(exp_avg_sq32, g32)]
+
+    copy_stochastic_list_(exp_avg, exp_avg32)
+    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
+    copy_stochastic_list_(grad, update)
 
-    # add the random number to the lower 16 bit of the mantissa
-    result.add_(source.view(dtype=torch.int32))
 
-    # mask off the lower 16 bit of the mantissa
+def adopt(grad, exp_avg_sq, exp_avg, beta1, beta2, step):
+    exp_avg, exp_avg_sq, grad = list_guard(exp_avg, exp_avg_sq, grad, y)
+    beta1, beta2, step = scalar_guard(beta1, beta2, step, lr, exp_avg[0])
+    _compilable_adopt_(grad, exp_avg_sq, exp_avg, beta1, beta2, step)
+    return grad
+
+
+def stochastic_round_list_(ref: List[Tensor], source: List[Tensor]):
+    return [stochastic_round_(r, s) for r, s in zip(ref, source)]
+
+
+@decorator_knowngood
+def stochastic_round_(ref: Tensor, source: Tensor):
+    if source.dtype == torch.bfloat16 or ref.dtype == source.dtype:
+        return source
+    if ref.dtype != torch.bfloat16:
+        return source.to(ref.dtype)
+    result = torch.randint_like(source, dtype=torch.int32, low=0, high=(1 << 16))
+    result.add_(source.view(dtype=torch.int32))
     result.bitwise_and_(-65536)  # -65536 = FFFF0000 as a signed int32
+    return result.view(dtype=torch.float32).bfloat16()
 
-    # copy the higher 16 bit into the target tensor
-    target.copy_(result.view(dtype=torch.float32))
+
+@decorator_knowngood
+def _compilable_copy_stochastic_(target: Tensor, source: Tensor):
+    target.copy_(stochastic_round_(target, source))
 
 
 def copy_stochastic_(target: Tensor, source: Tensor):
@@ -759,7 +920,7 @@ def _compilable_update_(p: List[Tensor], u: List[Tensor], decay: Tensor, add_fn:
 
 def update_param_(param: List[Tensor], update: List[Tensor], lr: float, decay: float, add_fn: callable = None,
                   caution: bool = False, grad: List[Tensor] = None):
-    param, update, grad = list_guard(param), list_guard(update), list_guard(grad)
+    param, update, grad = list_guard(param, update, grad)
     lr = scalar_guard(lr, param[0])
     if not caution:
         grad = [None] * len(param)
@@ -865,11 +1026,15 @@ def psgd_balance_Q(Q_in):
 
 
 def psgd_calc_A_and_conjB(exprA, G, Q):
+    V = torch.randn(G.shape, dtype=G.dtype, device=G.device)
+    eps = scalar_guard(math.sqrt(torch.finfo(torch.float32).eps), G)
+    eps *= G.norm() / G.numel()
+    G += V * eps
     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.randn(G.shape[1:] + G.shape[:1], dtype=promote(G.dtype), device=G.device)
+    conjB = torch.permute(V, p[1:] + p[:1]).to(promote(G.dtype))
     Q = [promote(q) for q in Q]
     for i, q in enumerate(Q):
         if q.dim() <= 1:
@@ -902,7 +1067,7 @@ def psgd_lb(A, max_abs):
 
 
 @decorator
-def psgd_update_precond(Q, exprs, G, precond_lr, tiny, oq, store_triu_as_line):
+def psgd_update_precond(Q, exprs, G, precond_lr, oq, store_triu_as_line):
     """Update Kronecker product preconditioner Q with pair (V, G)."""
     exprA, exprGs, _ = exprs
 
@@ -923,10 +1088,10 @@ def psgd_update_precond(Q, exprs, G, precond_lr, tiny, oq, store_triu_as_line):
         norm = term2.norm(float('inf'))
         if q.dim() < 2:
             term1 *= q.to(term1.dtype)
-            term1 /= norm.clamp_(min=tiny)
+            term1 /= norm.clamp_(min=tiny_bf16)
         else:
             torch.triu(term1, out=term1)
-            term1 /= psgd_lb(term2, norm).clamp_(tiny)
+            term1 /= psgd_lb(term2, norm).clamp_(tiny_bf16)
             torch.matmul(term1, q, out=term1)
         if store_triu_as_line:
             term1 = triu_to_line([term1])[0][1]
@@ -935,22 +1100,32 @@ def psgd_update_precond(Q, exprs, G, precond_lr, tiny, oq, store_triu_as_line):
 
 
 @decorator_knowngood
-def psgd_precond_grad(inplace: bool, exprs: str, grad: Tensor, *preconds: Tensor):
-    """Precondition gradient G with preconditioner Q."""
-    md = min_dtype(preconds)
-    out = torch.einsum(exprs, *[q.conj().to(md) for q in preconds], *[q.to(md) for q in preconds], grad.to(md))
-    if inplace:
-        set_(grad, out)
-        return grad
-    return out.to(grad.dtype)
+def _compilable_l2_clip_(x):
+    ref = x
+    x = list(map(promote, x))
+    norm = torch._foreach_norm(x)
+    torch._foreach_maximum_(norm, 1e-8)
+    out = torch._foreach_div(x, norm)
+    return stochastic_round_list_(ref, out)
+
 
+def l2_clip_(x):
+    x = list_guard(x)
+    return _compilable_l2_clip_(x)
 
-def norm_clip_(x, scale=None):
+
+@decorator_knowngood
+def _compilable_rmsnorm_clip_(x):
+    x = list(map(promote, x))
     norm = torch._foreach_norm(x)
-    if scale is not None:
-        torch._foreach_div_(norm, scale)
-    torch._foreach_div_(x, norm)
-    return x
+    norm = [n.div_(x_.numel() ** 0.5) for n, x_ in zip(norm, x)]
+    torch._foreach_maximum_(norm, 1e-6)
+    return torch._foreach_div(x, norm)
+
+
+def rmsnorm_clip_(x):
+    x = list_guard(x)
+    return _compilable_rmsnorm_clip_(x)
 
 
 def mu_law_compress(x, mu=127.0):
@@ -990,18 +1165,24 @@ def identity(x):
     return x
 
 
-def trust_region_clip_(grad, lerp: float = 0.9, scale: float = 1.5):
-    torch._foreach_mul_(grad, 1 / scale)
-    tanh = torch._foreach_tanh(grad)
-    torch._foreach_abs_(grad)
-    torch._foreach_log1p_(grad)
-    grad = [p.copysign_(t) for t, p in zip(tanh, grad)]  # torch doesn't have a foreach copysign
-    torch._foreach_lerp_(grad, tanh, lerp)  # sgn(x) * log(1 + |x|) * 0.1 + tanh(x) * 0.9
-    torch._foreach_mul_(grad, scale)
+@decorator_knowngood
+def _compilable_trust_region_clip_(grad, lerp: float = 0.9, scale: float = 1.5):
+    g32 = list(map(promote, grad))
+    [g.mul_(1 / scale) for g in g32]
+    tanh = torch._foreach_tanh(g32)
+    torch._foreach_abs_(g32)
+    torch._foreach_log1p_(g32)
+    [g.copysign_(t).lerp_(t, lerp).mul_(scale) for t, g in zip(tanh, g32)]
 
-    torch._foreach_maximum_(grad, -2)
-    torch._foreach_minimum_(grad, 2)
-    return grad
+    torch._foreach_maximum_(g32, -2)
+    torch._foreach_minimum_(g32, 2)
+    return [stochastic_round_(grad, g32) for grad, g32 in zip(grad, g32)]
+
+
+def trust_region_clip_(grad, lerp=0.9, scale=1.5):
+    grad = list_guard(grad)
+    lerp, scale = scalar_guard(lerp, scale, grad[0])
+    return _compilable_trust_region_clip_(grad, lerp, scale)
 
 
 @decorator
@@ -1040,60 +1221,57 @@ def update_triu_(q_state, materialised):
         copy_stochastic_(q, m)
 
 
-class PSGDBase(StatefulOptimizer):
-    balance_probability: float = 0.01
-
-    def __init__(self, parameters, groups, foreach: bool, stochastic_schedule: bool, clip_fn,
-                 preconditioner_update_probability):
-        super().__init__(parameters, {**groups, 'stochastic_schedule': stochastic_schedule}, foreach)
-        self.rng = random.Random(0x1923213)
-        self._tiny = torch.finfo(torch.bfloat16).tiny
-        if clip_fn is None:
-            clip_fn = identity
-        if preconditioner_update_probability is None:
-            preconditioner_update_probability = precond_update_prob_schedule()
-        self.clip_fn = clip_fn
-        self.preconditioner_update_probability = preconditioner_update_probability
-
-    def should_update(self, group, prob: Optional[float] = None, name: str = 'cumulative_prob'):
-        group[f'{name}_prob_step'] = group.get(f'{name}_prob_step', 0) + 1
-        if prob is None:
-            prob = self.preconditioner_update_probability(group[f'{name}_prob_step'])
-        if group['stochastic_schedule']:
-            return self.rng.random() < prob
-        cumulative_prob = group.get(name, 0)
-        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: 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):
-                if g.dim() > 1:
-                    if store_triu_as_line:
-                        psgd_balance_Q([q_ for _, q_ in q])
-                    else:
-                        psgd_balance_Q(q)
-
-
-# TODO: Figure out why this sometimes crashes
-# @decorator_knowngood
-def _compilable_precond_grad_cached_(ea: Tensor, expr: str, param: Tensor, lr: Tensor, weight_decay: Tensor,
-                                     clip_fn: callable, caution: bool, grad: Optional[Tensor], *cached_q: Tensor):
+def psgd_should_update(group, prob: Union[float, callable], rng: Optional[random.Random] = None,
+                       name: str = 'cumulative_prob'):
+    group[f'{name}_prob_step'] = group.get(f'{name}_prob_step', 0) + 1
+    if not isinstance(prob, float):
+        prob = prob(group[f'{name}_prob_step'])
+    if group['stochastic_schedule']:
+        return rng.random() < prob
+    cumulative_prob = state.get(name, 0)
+    group[name] = cumulative_prob + prob
+    return int(group[name]) > int(cumulative_prob)
+
+
+@decorator_knowngood
+def precond_grad_cached_(expr: str, ea: Tensor, *cached_q: Tensor, cast: bool = True):
     md = min_dtype(list(cached_q) + [ea])
     args = [q.to(md) for q in cached_q]
     args = args + [ea.to(md)]
     new = torch.einsum(expr, *args)
-    new = new.to(torch.float32)
-    _compilable_update_([param], clip_fn([new]), weight_decay, stochastic_add_, lr, caution, [grad])
+    if cast:
+        return new.to(ea.dtype)
+    return new
+
 
+@decorator_knowngood
+def _compilable_fused_precond_grad_cached_(expr: str, ea: Tensor, param, lr, grad, decay, caution, *cached_q: Tensor):
+    precond = precond_grad_cached_(expr, ea, *cached_q, cast=False)
+    update_param_(param, precond, lr, decay, caution=caution, grad=grad)
 
-def precond_grad_cached_(cached_q: List[Tensor], ea: Tensor, expr: str, param: Tensor, lr: float, weight_decay: float,
-                         clip_fn, caution, grad):
-    lr = scalar_guard(lr, param)
-    _compilable_precond_grad_cached_(ea, expr, param, lr, weight_decay, clip_fn, caution, grad, *cached_q)
+
+def fused_precond_grad_cached_(expr: str, ea: Tensor, param, lr, grad, decay, caution, *cached_q: Tensor):
+    lr = scalar_guard(lr, param[0])
+    _compilable_fused_precond_grad_cached_(expr, ea, param, lr, grad, decay, caution, *cached_q)
+
+
+@decorator_knowngood
+def psgd_precond_grad(expr: str, ea: Tensor, *preconds: Tensor):
+    md = min_dtype(list(preconds) + [ea])
+    args = [q.to(md) for q in preconds]
+    args = args + args + [ea.to(md)]
+    new = torch.einsum(expr, *args)
+    return new.to(ea.dtype)
+
+
+def _compilable_fused_psgd_precond_grad(expr: str, ea: Tensor, param, lr, grad, decay, caution, *preconds: Tensor):
+    precond = psgd_precond_grad(expr, grad, *preconds)
+    update_param_(param, precond, lr, decay, caution=caution, grad=grad)
+
+
+def fused_psgd_precond_grad(expr: str, ea: Tensor, param, lr, grad, decay, caution, *preconds: Tensor):
+    lr = scalar_guard(lr, param[0])
+    _compilable_fused_psgd_precond_grad(expr, ea, param, lr, grad, decay, caution, *preconds)
 
 
 @decorator_knowngood
@@ -1122,7 +1300,7 @@ def caution(g, update):
     _compilable_cautioning_(g, update)
 
 
-def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.001, flat_start=250):
+def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.001, flat_start=500):
     """Anneal preconditioner update probability during beginning of training.
 
     PSGD benefits from more preconditioner updates at the beginning of training,