heavyball 0.25.1__py3-none-any.whl → 1.0.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
@@ -16,6 +24,7 @@ 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
+tiny_bf16 = torch.finfo(torch.bfloat16).tiny
 
 
 def decorator(func):
@@ -60,41 +69,34 @@ 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:
         ckp1 = weight / weight_sum
     except ZeroDivisionError:
         ckp1 = 0
+    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)
+    lr, ckp1, beta1 = scalar_guard(lr, parameters[0]), scalar_guard(ckp1, parameters[0]), scalar_guard(beta1, parameters[0])
+    _compilable_schedule_free_(parameters, z, ckp1, grad, lr, beta1, decay)
     return weight_sum
 
 
@@ -162,10 +164,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
 
@@ -179,10 +184,27 @@ def exp_avg_sq_(state, grad, beta2, eps, out=None):
     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,
+                                     out: List[Optional[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)
+    return stochastic_round_list_(grad, out)
+
+
+def scale_by_exp_avg_sq_(grad, exp_avg_sq, beta2, eps):
+    grad, exp_avg_sq = list_guard(grad), list_guard(exp_avg_sq)
+    beta2, eps = scalar_guard(beta2, grad[0]), scalar_guard(eps, grad[0])
+    return _compilable_scale_by_exp_avg_sq_(grad, exp_avg_sq, beta2, eps, 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 +212,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), list_guard(gradients)
+    clip_val = scalar_guard(clip_val, parameters[0])
+    return _compilable_agc_(parameters, gradients, clip_val, minimum, eps)
 
 
 def is_compiling():
@@ -205,10 +236,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():
@@ -353,8 +381,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,6 +395,27 @@ 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):
+    beta = None
+    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)
     a = scalar_guard(a, x[0])
@@ -435,35 +482,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 +676,63 @@ 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)
-    return denom
+    return stochastic_round_list_(exp_avg, u32)
 
 
-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 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, exp_avg[0]), scalar_guard(beta2, exp_avg[0]), scalar_guard(step,
+                                                                                                        exp_avg[0])
+    return _compilable_adam_(exp_avg, exp_avg_sq, grad, beta1, beta2, step)
 
 
+@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)
+    _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 = map(list_guard, (y, exp_avg, exp_avg_sq, grad))
+    beta1, beta2, step, lr = [scalar_guard (x, y[0]) for x in (beta1, beta2, step, lr)]
+    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_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor],
-                         grad_projected: List[Tensor], beta1: Tensor, beta2: Tensor, step: Tensor):
+def _compilable_laprop_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad_projected: List[Tensor], beta1: Tensor,
+                        beta2: Tensor, step: Tensor):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -709,27 +745,109 @@ def _compilable_laprop_exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor],
     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):
+def laprop_(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)
+    _compilable_laprop_(exp_avg, exp_avg_sq, grad_projected, beta1, beta2, step)
+    return exp_avg
 
 
 @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_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)
 
-    # add the random number to the lower 16 bit of the mantissa
-    result.add_(source.view(dtype=torch.int32))
+    gp32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad_projected, exp_avg_sq]]
 
-    # mask off the lower 16 bit of the mantissa
+    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 fused_laprop_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad_projected: List[Tensor],
+                  beta1: float, beta2: float, step: int, lr: float, decay: float, caution: bool):
+    y, exp_avg, exp_avg_sq, grad_projected = map(list_guard, (y, exp_avg, exp_avg_sq, grad_projected))
+    beta1, beta2, step, lr = [scalar_guard (x, y[0]) for x in (beta1, beta2, step, lr)]
+    _fused_compilable_laprop_(y, exp_avg, exp_avg_sq, grad_projected, beta1, beta2, step, lr, decay, caution)
+
+
+@decorator_knowngood
+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):
+    y, grad, exp_avg_sq, exp_avg = list_guard(y), list_guard(grad), list_guard(exp_avg_sq), list_guard(exp_avg)
+    beta1, beta2, step, lr = [scalar_guard (x, y[0]) for x in (beta1, beta2, step, lr)]
+    _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)
+
+    return update
+
+
+def adopt(grad, exp_avg_sq, exp_avg, beta1, beta2, step):
+    grad, exp_avg_sq, exp_avg = list_guard(grad), list_guard(exp_avg_sq), list_guard(exp_avg)
+    beta1, beta2, step = scalar_guard(beta1, grad[0]), scalar_guard(beta2, grad[0]), scalar_guard(step, grad[0])
+    return _compilable_adopt_(grad, exp_avg_sq, exp_avg, beta1, beta2, step)
+
+
+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):
@@ -902,7 +1020,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 +1041,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 +1053,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 +1118,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_(g32, -2)
+    torch._foreach_minimum_(g32, 2)
+    return [stochastic_round_(grad, g32) for grad, g32 in zip(grad, g32)]
 
-    torch._foreach_maximum_(grad, -2)
-    torch._foreach_minimum_(grad, 2)
-    return grad
+
+def trust_region_clip_(grad, lerp=0.9, scale=1.5):
+    grad = list_guard(grad)
+    lerp, scale = scalar_guard(lerp, grad[0]), scalar_guard(scale, grad[0])
+    return _compilable_trust_region_clip_(grad, lerp, scale)
 
 
 @decorator
@@ -1040,60 +1174,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 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 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_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 +1253,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,