heavyball 1.6.2__py3-none-any.whl → 1.7.0__py3-none-any.whl

heavyball/utils.py

@@ -4,7 +4,7 @@ import math
 import random
 import string
 import warnings
-from typing import List, Optional, Tuple, Callable, Union
+from typing import Callable, List, Optional, Tuple, Union
 from unittest.mock import patch
 
 import numpy as np
@@ -17,25 +17,18 @@ from torch.utils._pytree import tree_map
 
 config.cache_size_limit = 2 ** 16
 
-np.warnings = warnings
-
 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
+zeroth_power_mode = "qr"  # 'qr' is baseline, 'newtonschulz' converges better and faster
 tiny_bf16 = torch.finfo(torch.bfloat16).tiny
 
-base_args = {'betas': (0.9, 0.999), 'precondition_frequency': 1, 'merge_dims': False, 'warmup_steps': 100,
-             'max_precond_dim': 2 ** 16, 'beta': 0.9, 'max_size_triangular': 2 ** 16, 'split': False, 'eps': 1e-8,
-             'weight_decay': 1e-4}
-
 
 def decorator(func):
     compiled = None
 
     @functools.wraps(func)
     def _fn(*args, **kwargs):
-        disable = compile_mode_recommended_to_none is None
         if is_compiling() or compile_mode_recommended_to_none is None:
             return func(*args, **kwargs)
         nonlocal compiled
@@ -66,7 +59,7 @@ einsum_base = string.ascii_lowercase
 
 @decorator_knowngood
 def _compilable_schedule_free_(p: List[Tensor], z: List[Tensor], ckp1: Tensor, update: List[Tensor], lr: Tensor,
-                               beta1: Tensor, decay: float, grad: List[Tensor], caution):
+        beta1: Tensor, decay: float, grad: List[Tensor], caution, ):
     for op, oz, u_, g_ in zip(p, z, update, grad):
         u_ = u_.view_as(op)
         p_, z_, u_ = map(promote, (op, oz, u_))
@@ -82,8 +75,8 @@ def _compilable_schedule_free_(p: List[Tensor], z: List[Tensor], ckp1: Tensor, u
 
 
 def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1: float, parameters: List[Tensor],
-                   z: List[Tensor], update: List[Tensor], grad: List[Tensor], caution: bool = False, r: float = 0.0,
-                   step: int = 0, decay: float = 0.0):
+        z: List[Tensor], update: List[Tensor], grad: List[Tensor], caution: bool = False, 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
 
@@ -165,7 +158,7 @@ def eps_sqrt(item, eps):
 
 @decorator_knowngood
 def _compilable_exp_avg_sq_(state: List[Tensor], grad: List[Tensor], beta2: Tensor, eps: Tensor,
-                            out: List[Optional[Tensor]]):
+        out: List[Optional[Tensor]]):
     g32 = promote(grad)
     s32 = _lerp(state, torch._foreach_mul(g32, g32), beta2)
 
@@ -227,7 +220,7 @@ def _compilable_agc_(parameters: List[Tensor], gradients: List[Tensor], clip_val
 
 
 def adaptive_gradient_clipping_(parameters: List[Tensor], gradients: List[Tensor], clip_val: float,
-                                minimum: float = 1e-3, eps: float = 1e-8):
+        minimum: float = 1e-3, eps: float = 1e-8):
     if clip_val <= 0:
         return gradients
     parameters, gradients = list_guard(parameters, gradients)
@@ -253,23 +246,22 @@ def clean():
 
 
 def _ignore_warning(msg):
-    warnings.filterwarnings('ignore', f'.*{msg}.*')
+    warnings.filterwarnings("ignore", f".*{msg}.*")
 
 
-def set_torch(benchmark_limit: int = 32):
+def set_torch(benchmark_limit: int = 32, einsum_strategy: str = "auto"):
     cudnn.benchmark = True
     cudnn.deterministic = False
     cudnn.benchmark_limit = benchmark_limit
     torch.use_deterministic_algorithms(False)
     torch.set_float32_matmul_precision("high")  # highest: FP32, high: TF32, medium: bf16
-    opt_einsum.enabled = False
-    opt_einsum.strategy = "auto"
+    opt_einsum.set_flags(True, einsum_strategy)
 
     # Torch calls these for 2nd-order optimization in HeavyBall, but they are explicitly handled.
     _ignore_warning(
-        'Using backward() with create_graph=True will create a reference cycle between the parameter and its gradient which can cause a memory leak')
+        "Using backward() with create_graph=True will create a reference cycle between the parameter and its gradient which can cause a memory leak")
     _ignore_warning(
-        'We recommend using autograd.grad when creating the graph to avoid this. If you have to use this function, make sure to reset the .grad fields of your parameters to None after use to break the cycle and avoid the leak')
+        "We recommend using autograd.grad when creating the graph to avoid this. If you have to use this function, make sure to reset the .grad fields of your parameters to None after use to break the cycle and avoid the leak")
 
 
 @decorator
@@ -277,7 +269,7 @@ def zeropower_via_newtonschulz5(G, steps=5, eps=1e-7):
     assert len(G.shape) == 2
     a, b, c = (3.4445, -4.7750, 2.0315)
     X = G.to(torch.bfloat16 if G.dtype != torch.float64 else G.dtype)  # Preserve float64 if present
-    X /= (X.norm() + eps)  # ensure top singular value <= 1
+    X /= X.norm() + eps  # ensure top singular value <= 1
     if G.size(0) > G.size(1):
         X = X.T
     for _ in range(steps):
@@ -290,10 +282,10 @@ def zeropower_via_newtonschulz5(G, steps=5, eps=1e-7):
 
 
 def ortho(x):
-    if zeroth_power_mode == 'qr':
+    if zeroth_power_mode == "qr":
         return torch.linalg.qr(x).Q
-    if zeroth_power_mode == 'svd':
-        u, s, v = torch.linalg.svd(x)
+    if zeroth_power_mode == "svd":
+        u, _s, v = torch.linalg.svd(x)
         return u @ v.T
     raise NotImplementedError(f"Unknown zeroth_power_mode: {zeroth_power_mode}")
 
@@ -351,12 +343,12 @@ def _compilable_grafting(magnitude, direction):
 
 @decorator_knowngood
 def inplace_orthogonal_(x: Tensor, mode: str, out: Tensor, scale_mode: str):
-    if mode == 'newtonschulz' or x.shape[0] != x.shape[1]:
+    if mode == "newtonschulz" or x.shape[0] != x.shape[1]:
         y = zeropower_via_newtonschulz5(x, 5)
-    elif mode == 'qr':
+    elif mode == "qr":
         y = torch.linalg.qr(promote(x)).Q
-    elif mode == 'svd':
-        u, s, v = torch.linalg.svd(promote(x))
+    elif mode == "svd":
+        u, _s, v = torch.linalg.svd(promote(x))
         y = u @ v.T
     else:
         raise NotImplementedError(f"Unknown zeroth_power_mode: {mode}")
@@ -403,7 +395,7 @@ def get_orthogonal_matrix_QR(GG: List[Tensor], Q: List[Tensor], exp_avg: Optiona
         q_old = promote(q.data)
 
         tmp = m @ q_old
-        est_eig = torch.einsum('ij,ij->j', q_old, tmp)
+        est_eig = torch.einsum("ij,ij->j", q_old, tmp)
         sort_idx = torch.argsort(est_eig, descending=True)
 
         tmp[:, sort_idx], _ = torch.linalg.qr(tmp[:, sort_idx])
@@ -415,19 +407,19 @@ def get_orthogonal_matrix_QR(GG: List[Tensor], Q: List[Tensor], exp_avg: Optiona
         return
 
     assert exp_avg.ndim < 13, "exp_avg.ndim must be less than 13"
-    in_str = einsum_base[:exp_avg.dim()]
-    out_str = einsum_base[exp_avg.dim():2 * exp_avg.dim()]
+    in_str = einsum_base[: exp_avg.dim()]
+    out_str = einsum_base[exp_avg.dim(): 2 * exp_avg.dim()]
 
     from_shampoo = ",".join([o + i for m, i, o in zip(Q, in_str, in_str.upper()) if m is not None])
     if not from_shampoo:
         return
 
-    to_shampoo = ','.join([i + o for m, i, o in zip(new_qs, in_str.upper(), out_str) if m is not None])
-    out_str = ''.join([o if o in to_shampoo else i for i, o in zip(in_str, out_str)])
+    to_shampoo = ",".join([i + o for m, i, o in zip(new_qs, in_str.upper(), out_str) if m is not None])
+    out_str = "".join([o if o in to_shampoo else i for i, o in zip(in_str, out_str)])
 
-    subscripts = f'{in_str},{from_shampoo},{to_shampoo}->{out_str}'
+    subscripts = f"{in_str},{from_shampoo},{to_shampoo}->{out_str}"
     exp_avg_new = torch.einsum(subscripts, exp_avg, *[q for q in Q if q is not None],
-                               *[q for q in new_qs if q is not None])
+        *[q for q in new_qs if q is not None])
     copy_stochastic_(exp_avg, exp_avg_new)
 
     for q, q_new in zip(Q, new_qs):
@@ -453,11 +445,11 @@ def get_orthogonal_matrix(mat, max_eps: float = 1e-3, min_eps: float = 1e-30):
         while True:
             try:
                 eye = torch.eye(m.shape[0], device=m.device, dtype=m.dtype)
-                eigval, eigvec = torch.linalg.eigh(m + eps * eye)
+                _eigval, eigvec = torch.linalg.eigh(m + eps * eye)
                 eigvec = eigvec.to(device=device, dtype=dtype)
                 break
             except torch.OutOfMemoryError:
-                if m.device.type == 'cpu':
+                if m.device.type == "cpu":
                     raise
                 else:
                     m = m.cpu()
@@ -489,21 +481,21 @@ def _compilable_stochastic_lerp_(x: List[Tensor], y: List[Tensor], a: Union[floa
 
 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" 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 'palm' in group and group['palm'] is True and 'beta2_scale' in group:
+    if "palm" in group and group["palm"] is True and "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]
+        return 1 - step ** -group["beta2_scale"]
+    if "betas" in group:
+        return group["betas"][1]
     raise ValueError("Beta2 not found in group.")
 
 
@@ -580,9 +572,9 @@ def update_ggt(grad, GG, max_precond_dim, precondition_1d, beta):
         if not isinstance(m, Tensor):
             continue
         b = einsum_base[idx]
-        g0 = einsum_base[:grad.dim()]
+        g0 = einsum_base[: grad.dim()]
         g1 = g0.replace(b, b.upper())
-        outer_product = torch.einsum(f'{g0},{g1}->{b + b.upper()}', grad, grad)
+        outer_product = torch.einsum(f"{g0},{g1}->{b + b.upper()}", grad, grad)
         stochastic_lerp_(m, outer_product, 1 - beta)
 
 
@@ -623,19 +615,19 @@ def init_preconditioner(grad, state, max_precond_dim, precondition_1d):
     """
     Initializes the preconditioner matrices (L and R in the paper).
     """
-    state['GG'] = []  # Will hold all the preconditioner matrices (L and R in the paper).
+    state["GG"] = []  # Will hold all the preconditioner matrices (L and R in the paper).
     if grad.numel() > 1 and (grad.ndim > 1 or precondition_1d):
         for sh in grad.shape:
             if sh > max_precond_dim or sh == 1:
                 # via @francois-rozet: https://github.com/HomebrewML/HeavyBall/commit/8b86be04967e2d095136d5603724f488f2d46592#diff-a430393dd0a6ee393944a9ed16416115c175de2414cf4a96e647197697f265e9R621
-                state['GG'].append(None)
+                state["GG"].append(None)
             else:
-                state['GG'].append(torch.zeros(sh, sh, device=grad.device, dtype=grad.dtype))
+                state["GG"].append(torch.zeros(sh, sh, device=grad.device, dtype=grad.dtype))
     else:
-        state['GG'].append(None)
+        state["GG"].append(None)
 
-    update_ggt(grad, state['GG'], max_precond_dim, precondition_1d, 0)
-    state['Q'] = get_orthogonal_matrix(state['GG'])
+    update_ggt(grad, state["GG"], max_precond_dim, precondition_1d, 0)
+    state["Q"] = get_orthogonal_matrix(state["GG"])
 
 
 @decorator
@@ -646,11 +638,11 @@ def project(grad, Q, back: bool):
     :param back: whether to project to Shampoo eigenbases or back to original space
     :return:
     """
-    param = einsum_base[:grad.dim()]
-    preconditioners = ",".join([(g + g.upper())[::-1 if back else 1] for m, g in zip(Q, param) if m is not None])
+    param = einsum_base[: grad.dim()]
+    preconditioners = ",".join([(g + g.upper())[:: -1 if back else 1] for m, g in zip(Q, param) if m is not None])
     if preconditioners:
-        out = ''.join([c.upper() if c.upper() in preconditioners else c for c in param])
-        out = torch.einsum(f'{param},{preconditioners}->{out}', promote(grad), *[q for q in Q if q is not None])
+        out = "".join([c.upper() if c.upper() in preconditioners else c for c in param])
+        out = torch.einsum(f"{param},{preconditioners}->{out}", promote(grad), *[q for q in Q if q is not None])
         grad = out.to(grad.dtype)
     return grad
 
@@ -667,12 +659,12 @@ def modify_closure(closure):
     """
 
     def patched_backward(self, *args, **kwargs):
-        kwargs['create_graph'] = True
+        kwargs["create_graph"] = True
         return original_backward(self, *args, **kwargs)
 
     original_backward = torch.Tensor.backward
 
-    with patch.object(torch.Tensor, 'backward', patched_backward):
+    with patch.object(torch.Tensor, "backward", patched_backward):
         return closure()
 
 
@@ -683,6 +675,7 @@ class StatefulOptimizer(torch.optim.Optimizer):
     The previous (heavyball<=1.5.3) default was `True`, which is incompatible with some benchmarks but works better with RevNet
     Further notice that both methods have different numerics outputs
     """
+
     ema_decay: float = 0.001
     compile_step: bool = False
     hessian_approx: bool = False
@@ -691,10 +684,10 @@ class StatefulOptimizer(torch.optim.Optimizer):
     finite_differences: bool = False
 
     def __init__(self, params, defaults, foreach: bool = True, use_ema: bool = False):
-        super().__init__(params, {**defaults, 'foreach': foreach})
+        super().__init__(params, {**defaults, "foreach": foreach})
         self.use_ema = use_ema
         self.mapping = {}
-        self._inner_group = {'stochastic_schedule': self.stochastic_schedule}
+        self._inner_group = {"stochastic_schedule": self.stochastic_schedule}
         self._precond_rng = random.Random(0x12312)
         self._is_preconditioning = None
 
@@ -710,24 +703,24 @@ class StatefulOptimizer(torch.optim.Optimizer):
     def mars_correct_list(self, group, p_list, g_list, mars_gamma, beta):
         for p, g in zip(p_list, g_list):
             state = self.state_(p)
-            if 'mars_old_grad' not in state:
-                state['mars_old_grad'] = torch.zeros_like(g)
-        old_gs = [self.state_(p)['mars_old_grad'] for p in p_list]
+            if "mars_old_grad" not in state:
+                state["mars_old_grad"] = torch.zeros_like(g)
+        old_gs = [self.state_(p)["mars_old_grad"] for p in p_list]
         mars_correction(g_list, old_gs, mars_gamma, beta)
 
     def split_p_and_g_in_group(self, group: dict, skip_none: bool = True, should_promote: bool = True,
-                               beta1: float = -1.0):
+            beta1: float = -1.0):
         for p in group["params"]:
             if p in self.mapping:
                 p_views = self.mapping[p]
             else:
                 self.mapping[p] = p_views = merge_group(group, p)
 
-            grad = getattr(p, 'grad', None)
+            grad = getattr(p, "grad", None)
             p.grad = None
 
             if grad is None:
-                grad = [getattr(pv, 'grad', None) for pv in p_views]
+                grad = [getattr(pv, "grad", None) for pv in p_views]
             else:
                 grad = merge_group(group, grad)
 
@@ -736,8 +729,8 @@ class StatefulOptimizer(torch.optim.Optimizer):
                     continue
                 if should_promote:
                     g = promote(g)
-                if beta1 >= 0 and group.get('mars', False):
-                    self.mars_correct_list(group, [pv], [g], group['mars_gamma'], beta1)
+                if beta1 >= 0 and group.get("mars", False):
+                    self.mars_correct_list(group, [pv], [g], group["mars_gamma"], beta1)
                 yield pv, g
 
     def state_size(self) -> int:
@@ -759,46 +752,46 @@ class StatefulOptimizer(torch.optim.Optimizer):
     def ema_update(self):
         with torch.no_grad():
             for group in self.param_groups:
-                active_p = [p for p in group['params']]
+                active_p = [p for p in group["params"]]
 
                 if not active_p:
                     return
 
-                k = group['ema_step'] = group.get('ema_step', -1) + 1
+                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)
+                    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])
+                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 group in self.param_groups:
-                active_p = [p for p in group['params']]
+                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):
+                    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)
+                        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 group in self.param_groups:
-                active_p = [p for p in group['params']]
+                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)
+                    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 _handle_closure(self, closure):
@@ -824,7 +817,8 @@ class StatefulOptimizer(torch.optim.Optimizer):
                     grads.append(g)
                     p.vector = torch.randn_like(p)
                     p.orig = p.data.clone()
-                    stochastic_add_(p.data, p.vector, tiny_bf16)
+                    # scale taken from https://github.com/lixilinx/psgd_torch/blob/1943e66596111e78157ca1b72b31c1dfdf0653ef/preconditioned_stochastic_gradient_descent.py#L2161
+                    stochastic_add_(p.data, p.vector, torch.finfo(p.dtype).eps ** 0.5)
         else:
             with torch.enable_grad():
                 loss = modify_closure(closure)
@@ -833,6 +827,8 @@ class StatefulOptimizer(torch.optim.Optimizer):
             with torch.enable_grad():
                 closure()
 
+            # we don't subtract the vector here again to avoid accumulating error from (x + eps - eps + eps - eps)
+            # this costs more memory, but the imprecision seems too severe to use the other method
             for group in self.param_groups:
                 for p, g in self.split_p_and_g_in_group(group, skip_none=True, should_promote=False):
                     p.grad = grads.pop(0)
@@ -845,7 +841,7 @@ class StatefulOptimizer(torch.optim.Optimizer):
                 for p, g in self.split_p_and_g_in_group(group, skip_none=True, should_promote=False):
                     p.grad = g
             params, grads = zip(*[x for group in self.param_groups for x in
-                                  self.split_p_and_g_in_group(group, skip_none=True, should_promote=False)])
+                self.split_p_and_g_in_group(group, skip_none=True, should_promote=False)])
             vs = [torch.randn_like(p) for p in params]
             with torch.enable_grad():
                 hvs = torch.autograd.grad(grads, params, vs)
@@ -867,7 +863,7 @@ class StatefulOptimizer(torch.optim.Optimizer):
         # we assume that parameters are constant and that there are no excessive recompiles
         with torch.no_grad(), torch._dynamo.utils.disable_cache_limit():
             for group in self.param_groups:
-                group['is_preconditioning'] = self._is_preconditioning
+                group["is_preconditioning"] = self._is_preconditioning
                 self._step(group)
                 if self.use_ema:
                     self.ema_update()
@@ -892,7 +888,7 @@ def _lerp(state: List[Tensor], grad: List[Tensor], beta):
 
 @decorator_knowngood
 def _compilable_adam_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: Tensor, beta2: Tensor,
-                      step: Tensor, eps: Tensor):
+        step: Tensor, eps: Tensor, ):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -904,7 +900,7 @@ def _compilable_adam_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: Lis
 
 
 def adam_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: float, beta2: float, step: int,
-          eps: float = 1e-8):
+        eps: float = 1e-8, ):
     exp_avg, exp_avg_sq, grad = map(list_guard, (exp_avg, exp_avg_sq, grad))
     beta1, beta2, step, eps = scalar_guard(beta1, beta2, step, eps, exp_avg[0])
     _compilable_adam_(exp_avg, exp_avg_sq, grad, beta1, beta2, step, eps)
@@ -913,8 +909,8 @@ def adam_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], b
 
 @decorator_knowngood
 def _fused_compilable_adam_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor], update: List[Tensor],
-                            grad: List[Tensor], beta1: Tensor, beta2: Tensor, step: Tensor, decay: Tensor, lr: Tensor,
-                            eps: Tensor, caution: bool):
+        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)
 
@@ -926,8 +922,8 @@ def _fused_compilable_adam_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq:
 
 
 def fused_adam_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor], update: List[Tensor],
-                grad: List[Tensor], beta1: float, beta2: float, step: int, lr: float, eps: float, decay: float,
-                caution: bool):
+        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])
     _fused_compilable_adam_(y, exp_avg, exp_avg_sq, update, grad, beta1, beta2, step, decay, lr, eps, caution)
@@ -935,7 +931,7 @@ def fused_adam_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor]
 
 @decorator_knowngood
 def _compilable_laprop_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: Tensor,
-                        beta2: Tensor, step: Tensor, eps: Tensor):
+        beta2: Tensor, step: Tensor, eps: Tensor, ):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -947,7 +943,7 @@ def _compilable_laprop_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: L
 
 
 def laprop_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: float, beta2: float, step: int,
-            eps: float = 1e-8):
+        eps: float = 1e-8, ):
     exp_avg, exp_avg_sq, grad = list_guard(exp_avg, exp_avg_sq, grad)
     beta1, beta2, step, eps = scalar_guard(beta1, beta2, step, eps, exp_avg[0])
     _compilable_laprop_(exp_avg, exp_avg_sq, grad, beta1, beta2, step, eps)
@@ -956,8 +952,8 @@ def laprop_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor],
 
 @decorator_knowngood
 def _fused_compilable_laprop_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor], update: List[Tensor],
-                              grad: List[Tensor], beta1: Tensor, beta2: Tensor, step: Tensor, lr: Tensor, decay: Tensor,
-                              caution: bool, eps: Tensor):
+        grad: List[Tensor], beta1: Tensor, beta2: Tensor, step: Tensor, lr: Tensor, decay: Tensor, caution: bool,
+        eps: Tensor, ):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -969,8 +965,8 @@ def _fused_compilable_laprop_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq
 
 
 def fused_laprop_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tensor], update: List[Tensor],
-                  grad: List[Tensor], beta1: float, beta2: float, step: int, lr: float, decay: float, caution: bool,
-                  eps: float = 1e-8):
+        grad: List[Tensor], beta1: float, beta2: float, step: int, lr: float, decay: float, caution: bool,
+        eps: float = 1e-8, ):
     exp_avg, exp_avg_sq, grad, y = list_guard(exp_avg, exp_avg_sq, grad, y)
     beta1, beta2, step, lr, eps = scalar_guard(beta1, beta2, step, lr, eps, exp_avg[0])
     _fused_compilable_laprop_(y, exp_avg, exp_avg_sq, update, grad, beta1, beta2, step, lr, decay, caution, eps)
@@ -978,7 +974,7 @@ def fused_laprop_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq: List[Tenso
 
 @decorator_knowngood
 def _fused_compilable_adopt_(y, update, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay, caution):
-    u32, g32, exp_avg_sq32, exp_avg32 = [list(map(promote, x)) for x in [update, grad, exp_avg_sq, exp_avg]]
+    u32, g32, exp_avg_sq32 = [list(map(promote, x)) for x in [update, grad, exp_avg_sq]]
     _compilable_update_(y, u32, decay, lr, caution, g32)
 
     beta1 = beta_debias(beta1, step)
@@ -997,7 +993,7 @@ def fused_adopt_(y, update, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, e
 
 @decorator_knowngood
 def _compilable_adopt_(grad, exp_avg_sq, exp_avg, beta1, beta2, step, eps):
-    g32, exp_avg32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg, exp_avg_sq]]
+    g32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg_sq]]
     update = [e.clone() for e in exp_avg]
 
     beta1 = beta_debias(beta1, step)
@@ -1045,7 +1041,7 @@ def copy_stochastic_(target: Tensor, source: Tensor):
 
 @decorator_knowngood
 def _compilable_update_(p: List[Tensor], u: List[Tensor], decay: Tensor, lr: Tensor, caution: bool,
-                        g: List[Optional[Tensor]]):
+        g: List[Optional[Tensor]]):
     for u_, g_, p_ in zip(u, g, p):  # lr is data-dependent -> can't compile a foreach
         u_ = promote(u_.view_as(p_))
         p32_ = promote(p_)
@@ -1056,7 +1052,7 @@ def _compilable_update_(p: List[Tensor], u: List[Tensor], decay: Tensor, lr: Ten
 
 
 def update_param_(param: List[Tensor], update: List[Tensor], lr: float, decay: float, caution: bool = False,
-                  grad: List[Tensor] = None):
+        grad: List[Tensor] = None):
     param, update, grad = list_guard(param, update, grad)
     lr = scalar_guard(lr, param[0])
     if not caution:
@@ -1064,38 +1060,70 @@ def update_param_(param: List[Tensor], update: List[Tensor], lr: float, decay: f
     _compilable_update_(param, update, decay, lr, caution, grad)
 
 
-def precond_schedule(step, precond_scheduler, rng):
+def precond_schedule(step, precond_scheduler):
     precond_prob = max(step, 1) ** precond_scheduler[0]
     precond_prob = math.log10(precond_prob)
     precond_prob = precond_prob ** precond_scheduler[1] + 1
-    precond_prob = 1 / precond_prob
-    update_precond = rng.random() < precond_prob
-    return update_precond
+    return 1 / precond_prob
 
 
 def get_soap_precond_schedule(precond_scheduler):
-    rng = random.Random(0x12312)
-
-    def _inner(step):
-        return precond_schedule(step, precond_scheduler, rng)
-
-    return _inner
+    return functools.partial(precond_schedule, precond_scheduler=precond_scheduler)
 
 
 def _max_idx(x: List[int]):
     return len(x) - 1 - np.argmax(x[::-1])  # we want to start counting from the back, as torch is fan-out/fan-in
 
 
-def init_Q_exprs(t, scale, max_size, min_ndim_triangular, memory_save_mode, dtype=None):
-    """For a scalar or tensor t, we initialize its preconditioner Q and
+@decorator_knowngood
+def mean_root(x: torch.Tensor, pow: float):
+    return stochastic_round_(x, x.float().pow(pow).mean().pow(-1 / pow / 2))
+
+
+@decorator_knowngood
+def divided_root(x, y, pow0, pow1):
+    mean_x = x.float().pow(pow0).mean().pow(1 / pow0 / 2)
+    mean_y = y.float().pow(pow1).mean().pow(-1 / pow1 / 2)
+    return stochastic_round_(x, mean_x * mean_y)  # multiply here, as we already divide in pow -1
+
+
+def precond_init_scale(scale, scale_scale, grad, hessian_vector, vector):
+    if scale is not None:
+        warn_once(
+            "It's recommended to use precond_init_scale=None (default since 1.7.x), which uses advanced heuristics.")
+        if scale_scale is not None:
+            warn_once(
+                "precond_init_scale_scale multiplies the precond_init_scale by a constant factor. With a fixed precond_init_scale, you should explicitly multiply it into the precond_init_scale.")
+        return scale
+    if hessian_vector is None:
+        return mean_root(grad, 4) * scale_scale
+    return divided_root(vector, hessian_vector, 2, 4) * scale_scale
+
+
+def init_lra(grad, scale, scale_scale, rank, hessian_vector, vector, dtype=None):
+    scale = precond_init_scale(scale, scale_scale, grad, hessian_vector, vector)
+    U = torch.randn((*grad.shape, rank), dtype=dtype, device=grad.device)
+    V = torch.randn((*grad.shape, rank), dtype=dtype, device=grad.device)
+    d = torch.full_like(grad, scale, dtype=dtype, device=grad.device)
+    return U, V, d
+
+
+def init_Q_exprs(grad, scale, scale_scale, max_size, min_ndim_triangular, memory_save_mode, hessian_vector, vector,
+        dtype=None):
+    """
+    For a scalar or tensor `grad`, we initialize its preconditioner Q and
     reusable einsum expressions for updating Q and preconditioning gradient.
+
+    precond init scale computation from
+    https://github.com/lixilinx/psgd_torch/blob/1943e66596111e78157ca1b72b31c1dfdf0653ef/preconditioned_stochastic_gradient_descent.py#L2208-L2227
     """
+    scale = precond_init_scale(scale, scale_scale, grad, hessian_vector, vector)
     letters = string.ascii_lowercase + string.ascii_uppercase
-    dtype = dtype if dtype is not None else t.dtype
-    shape = t.shape
+    dtype = dtype if dtype is not None else grad.dtype
+    shape = grad.shape
 
     if len(shape) == 0:  # scalar
-        Q = [scale * torch.ones_like(t, dtype=dtype)]
+        Q = [scale * torch.ones_like(grad, dtype=dtype)]
         exprA = ",->"
         exprGs = [",->"]
         exprP = ",,->"
@@ -1103,7 +1131,7 @@ def init_Q_exprs(t, scale, max_size, min_ndim_triangular, memory_save_mode, dtyp
 
     # Tensor
     if len(shape) > 13:
-        raise ValueError(f"Got tensor with dim {len(t.shape)}; Einstein runs out of letters!")
+        raise ValueError(f"Got tensor with dim {len(grad.shape)}; Einstein runs out of letters!")
 
     scale = scale ** (1 / len(shape))
 
@@ -1129,7 +1157,7 @@ def init_Q_exprs(t, scale, max_size, min_ndim_triangular, memory_save_mode, dtyp
     for i, (size, dim_d) in enumerate(zip(shape, dim_diag)):
         if size == 1 or size > max_size or len(shape) < min_ndim_triangular or dim_d:
             # use diagonal matrix as preconditioner for this dim
-            Q.append(scale * torch.ones(size, dtype=promote(dtype), device=t.device))
+            Q.append(scale * torch.ones(size, dtype=promote(dtype), device=grad.device))
 
             piece1A.append(letters[i])
             piece2A = piece2A + letters[i]
@@ -1143,13 +1171,13 @@ def init_Q_exprs(t, scale, max_size, min_ndim_triangular, memory_save_mode, dtyp
             piece4P = piece4P + letters[i + 13]
         else:
             # use triangular matrix as preconditioner for this dim
-            Q.append(scale * torch.eye(size, dtype=dtype, device=t.device))
+            Q.append(scale * torch.eye(size, dtype=dtype, device=grad.device))
             piece1A.append(letters[i] + letters[i + 13])
             piece2A = piece2A + letters[i + 13]
             piece3A = piece3A + letters[i]
             piece1 = "".join([(letters[i + 13] if j == i else letters[j]) for j in range(len(shape))])
             piece2 = "".join([(letters[i + 26] if j == i else letters[j]) for j in range(len(shape))])
-            subscripts = (piece1 + "," + piece2 + "->" + letters[i + 13] + letters[i + 26])
+            subscripts = piece1 + "," + piece2 + "->" + letters[i + 13] + letters[i + 26]
             exprGs.append(subscripts)
             a, b, c = (letters[i], letters[i + 13], letters[i + 26])
             piece1P.append(a + b)
@@ -1158,7 +1186,7 @@ def init_Q_exprs(t, scale, max_size, min_ndim_triangular, memory_save_mode, dtyp
             piece4P = piece4P + b
 
     exprA = ",".join(piece1A) + "," + piece2A + "->" + piece3A
-    exprP = (",".join(piece1P) + "," + ",".join(piece2P) + "," + piece3P + "->" + piece4P)
+    exprP = ",".join(piece1P) + "," + ",".join(piece2P) + "," + piece3P + "->" + piece4P
     return [Q, (exprA, tuple(exprGs), exprP)]
 
 
@@ -1170,17 +1198,171 @@ def psgd_balance_Q(Q_in):
     torch._foreach_mul_(Q_in, list(norms))
 
 
+@decorator
+def psgd_balance_lra(U: Tensor, V: Tensor):
+    u_norm = promote(torch.linalg.vector_norm(U))
+    v_norm = promote(torch.linalg.vector_norm(V))
+    scale = (u_norm / v_norm) ** 0.5
+    U.div_(scale)
+    V.mul_(scale)
+
+
+@decorator
+def low_rank_mm(U: Tensor, V: Tensor, x: Tensor) -> Tensor:
+    dtype = min_dtype([U, V, x])
+    return x + torch.einsum("br,gr,g->b", U.to(dtype), V.to(dtype), x.to(dtype)).to(x.dtype)
+
+
+def update_lra_precond_(U: List[Tensor], V: List[Tensor], d: List[Tensor], vector: Tensor, hessian_vector: Tensor,
+        eps: float, step: float, delayed: bool, ):
+    """
+    Adapted from https://github.com/lixilinx/psgd_torch/blob/6dbea94915679d08a289928e6431b6ce07931aaf/preconditioned_stochastic_gradient_descent.py#L657
+    """
+    U_orig, V_orig, d_orig = U, V, d
+
+    U, V, d = flatten(U, 1), flatten(V, 1), flatten(d)
+
+    dtype = min_dtype([U, V, vector, hessian_vector])
+    U, V, vector, hessian_vector = U.to(dtype), V.to(dtype), vector.to(dtype), hessian_vector.to(dtype)
+
+    eps = scalar_guard(eps, vector)
+
+    Qh = low_rank_mm(U, V, d * hessian_vector)
+    Ph = d * low_rank_mm(V, U, Qh)
+    rank = U.size(1)
+
+    VtU = torch.einsum("br,bn->rn", V, U)  # (rank, rank)
+    I = torch.eye(rank, dtype=VtU.dtype, device=VtU.device)
+    IpVtU = I + VtU
+    invQtv = vector / d
+
+    # LU factorization to reuse computation
+    try:
+        LU, pivots = torch.linalg.lu_factor(IpVtU)
+    except RuntimeError:
+        # Error:
+        # U[2,2] is zero and using it on lu_solve would result in a division by zero.
+        # If you still want to perform the factorization, consider calling
+        # linalg.lu(A, pivot) or linalg.lu_factor_ex(A, pivot)
+        # ---
+        # So, we skip this step and reattempt on the next one
+        return U.to(U_orig[0].dtype), V.to(V_orig[0].dtype), d.to(d_orig[0].dtype)
+
+    invQtv = invQtv - V @ torch.linalg.lu_solve(LU, pivots, (U.T @ invQtv).view(-1, 1), adjoint=True).flatten()
+    invPv = invQtv - U @ torch.linalg.lu_solve(LU, pivots, (V.T @ invQtv).view(-1, 1)).flatten()
+    invPv = invPv / d
+
+    nablaD = Ph * hessian_vector - vector * invPv
+    divisor = (Ph.square() + vector.square()) * (hessian_vector.square() + invPv.square())
+    divisor = divisor.add(eps).sqrt().max()
+    d_step = step / divisor
+
+    apply_flat_add(d_orig, d * nablaD, -d_step)
+
+    a, b = Qh, invQtv
+
+    precond_u = random.random() < 0.5  # update either U or V, not both at the same time
+    precond = V if precond_u else U
+    atV = torch.einsum("b,br->r", a, precond)  # o == one
+    btV = torch.einsum("b,br->r", b, precond)
+    atVVt = torch.einsum("r,br->b", atV, precond)
+    btVVt = torch.einsum("r,br->b", btV, precond)
+    precond_step = step / (a.norm() * atVVt.norm() + b.norm() * btVVt.norm() + eps)
+    if precond_u:
+        a = torch.einsum("b,r,rg->bg", a, atV, IpVtU)
+        b = torch.einsum("b,r,rg->bg", b, btV, IpVtU)
+    else:
+        a = a + torch.einsum("br,r->b", V, atV)
+        b = b + torch.einsum("br,r->b", V, btV)
+        a = torch.einsum("b,r->br", a, atV)
+        b = torch.einsum("b,r->br", b, btV)
+    apply_flat_add(U_orig if precond_u else V_orig, b - a, precond_step)
+
+    if not delayed:
+        stochastic_add_([d], [d * nablaD], -d_step)
+        stochastic_add_([U if precond_u else V], [b - a], precond_step)
+    return U.to(U_orig[0].dtype), V.to(V_orig[0].dtype), d.to(d_orig[0].dtype)
+
+
+def lra_precond(U, V, d, g):
+    """
+    As-is from https://github.com/lixilinx/psgd_torch/blob/6dbea94915679d08a289928e6431b6ce07931aaf/preconditioned_stochastic_gradient_descent.py#L744
+    """
+    g = low_rank_mm(U, V, d * g)
+    return d * low_rank_mm(V, U, g)
+
+
+@decorator_knowngood
+def dampen_grad(g: Tensor, damp: float = 2 ** -13):
+    # https://github.com/lixilinx/psgd_torch/blob/1943e66596111e78157ca1b72b31c1dfdf0653ef/preconditioned_stochastic_gradient_descent.py#L50
+    v = torch.randn_like(g)
+    return v, g + damp * g.abs().mean() * v
+
+
+@decorator_knowngood
+def apply_lra_update(params: List[Tensor], update: Tensor, U: Tensor, V: Tensor, d: Tensor):
+    update = lra_precond(U, V, d, update)
+    start = 0
+    update = update.flatten()
+    for p in params:
+        size = p.numel()
+        copy_stochastic_(p, update[start: start + size].view_as(p))
+        start += size
+
+
+@decorator_knowngood
+def apply_flat_update(params: List[Tensor], update: Tensor):
+    start = 0
+    update = update.flatten()
+    for p in params:
+        size = p.numel()
+        copy_stochastic_(p, update[start: start + size].view_as(p))
+        start += size
+
+
+@decorator_knowngood
+def apply_flat_add(params: List[Tensor], update: Tensor, alpha: Tensor):
+    start = 0
+    update = update.flatten()
+    for p in params:
+        size = p.numel()
+        stochastic_add_([p], [update[start: start + size].view_as(p)], alpha)
+        start += size
+
+
+@decorator_knowngood
+def extract_from_flat_update(params: List[Tensor], update: Tensor):
+    start = 0
+    outputs = []
+    update = update.flatten()
+    for p in params:
+        size = p.numel()
+        outputs.append(update[start: start + size].view_as(p))
+        start += size
+    return outputs
+
+
+def flatten(x: List[Tensor], remaining: int = 0) -> Tensor:
+    return torch.cat([i.flatten(0, -1 - remaining) for i in x], 0)
+
+
+def dampen_multiple(g: List[Tensor], damp: float = 2 ** -13):
+    vs = []
+    gs = []
+    for g_ in g:
+        v, g = dampen_grad(g_, damp)
+        vs.append(v)
+        gs.append(g)
+    return flatten(vs), flatten(gs)
+
+
 def psgd_calc_A_and_conjB(exprA, G, Q, V=None):
-    eps = scalar_guard(math.sqrt(torch.finfo(G.dtype).eps), G)
-    eps *= G.norm() / G.numel()
-    G = G + torch.randn_like(G) * 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()
     if V is None:
-        conjB = torch.randn(G.shape[1:] + G.shape[:1], dtype=promote(G.dtype), device=G.device)
-    else:
-        conjB = V.permute(*range(1, order), 0).to(promote(G.dtype))
+        V, G = dampen_grad(G)
+    conjB = V.permute(*range(1, order), 0).to(promote(G.dtype))
+    md = min_dtype(Q + [G])
+    A = torch.einsum(exprA, *[q.to(md) for q in Q], G.to(md)).to(G.dtype)
     Q = [promote(q) for q in Q]
     for i, q in enumerate(Q):
         if q.dim() <= 1:
@@ -1195,12 +1377,12 @@ def psgd_calc_A_and_conjB(exprA, G, Q, V=None):
 
 def psgd_lb(A, max_abs):
     A /= max_abs
-    a0 = torch.einsum('ij,ij->j', A, A)
+    a0 = torch.einsum("ij,ij->j", A, A)
     i = torch.argmax(a0)
     x = torch.index_select(A, 1, i).flatten().contiguous()
-    x = torch.einsum('i,ij->j', x, A)
+    x = torch.einsum("i,ij->j", x, A)
     x /= x.norm()
-    x = torch.einsum('j,kj->k', x, A)
+    x = torch.einsum("j,kj->k", x, A)
     x = x.norm()
     x *= max_abs
     return x
@@ -1217,7 +1399,7 @@ def psgd_update_precond(Q, exprs, G, precond_lr, oq, store_triu_as_line, V):
         term2 = promote(torch.einsum(exprG, conjB, conjB))
         term1, term2 = term1 - term2, term1 + term2
         term1 *= precond_lr
-        norm = term2.norm(float('inf'))
+        norm = term2.norm(float("inf"))
         if q.dim() < 2:
             term1 *= q.to(term1.dtype) / norm.clamp_(min=tiny_bf16)
         else:
@@ -1245,7 +1427,7 @@ def l2_normalization_(x, clip_at: float = 1e-8):
     return _compilable_l2_clip_(x, clip_at)
 
 
-def l2_clip_(x, clip_at: float = 1.):
+def l2_clip_(x, clip_at: float = 1.0):
     x = list_guard(x)
     return _compilable_l2_clip_(x, clip_at)
 
@@ -1438,11 +1620,11 @@ def warn_once(msg):
 
 
 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
+        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']:
+        prob = prob(group[f"{name}_prob_step"])
+    if group["stochastic_schedule"]:
         return rng.random() < prob
     cumulative_prob = group.get(name, 0)
     group[name] = cumulative_prob + prob
@@ -1451,7 +1633,7 @@ def psgd_should_update(group, prob: Union[float, callable], rng: Optional[random
 
 @decorator_knowngood
 def precond_grad_cached_(expr: str, ea: Tensor, *cached_q: Tensor, caution: bool = False, grad: Optional[Tensor] = None,
-                         cast: bool = True):
+        cast: bool = True):
     if caution:
         ea = _compilable_cautioning(grad, ea)
     md = min_dtype(list(cached_q) + [ea])
@@ -1564,15 +1746,16 @@ def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.999, flat_
 
 
 def merge_group(group, *tensors):
-    if not group.get('merge_dims', False):
+    if not group.get("merge_dims", False):
         return tensors
     if isinstance(tensors[0], list):
         return [merge_group(group, *t) for t in tensors]
 
     out = []
     for t in tensors:
-        append_or_extend(out, dim_merger(t, group['max_size_triangular'] if 'max_size_triangular' in group else group[
-            'max_precond_dim'], group.get('split', False)))
+        append_or_extend(out,
+            dim_merger(t, group["max_size_triangular"] if "max_size_triangular" in group else group["max_precond_dim"],
+                group.get("split", False), ), )
     return out
 
 
@@ -1599,7 +1782,7 @@ def fused_hook(parameters, optimizer, *args, **kwargs):
     o = optimizer(parameters, *args, **kwargs)
     step_fn = o.step
     o.step = functools.partial(warn_once,
-                               msg="You're trying to call `step` on a fused optimizer. This will not do anything.")
+        msg="You're trying to call `step` on a fused optimizer. This will not do anything.")
 
     def _step(p: Tensor):
         seen_params.add(p)