heavyball 1.6.3__py3-none-any.whl → 1.7.1__py3-none-any.whl

heavyball/utils.py

@@ -1,11 +1,13 @@
+import contextlib
 import functools
 import gc
+import inspect
 import math
 import random
+import re
 import string
 import warnings
-from typing import List, Optional, Tuple, Callable, Union
-from unittest.mock import patch
+from typing import Callable, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -15,19 +17,22 @@ from torch._dynamo.exc import TorchDynamoException
 from torch.backends import cudnn, opt_einsum
 from torch.utils._pytree import tree_map
 
-config.cache_size_limit = 2 ** 16
-
-np.warnings = warnings
+config.cache_size_limit = 2**16
 
 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}
+_cudnn_double_backward_pattern = re.compile(
+    r"the derivative for .* is not implemented\. Double backwards .* To run double backwards"
+)
+_torch_compile_double_backward_pattern = re.compile(r"compile.*does not currently support double backward")
+_fd_error = (
+    "You can accelerate startup by globally enabling finite_differences first "  #
+    "(via opt.finite_differences=True or by subclassing it)\n"
+    "Original Error: "
+)
 
 
 def decorator(func):
@@ -35,7 +40,6 @@ def decorator(func):
 
     @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
@@ -65,8 +69,17 @@ 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):
+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,
+):
     for op, oz, u_, g_ in zip(p, z, update, grad):
         u_ = u_.view_as(op)
         p_, z_, u_ = map(promote, (op, oz, u_))
@@ -81,9 +94,20 @@ def _compilable_schedule_free_(p: List[Tensor], z: List[Tensor], ckp1: Tensor, u
         copy_stochastic_(oz, z_)
 
 
-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):
+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,
+):
     weight = abs(lr) ** weight_lr_power * max(step, 1) ** r
     weight_sum = weight_sum + weight
 
@@ -156,7 +180,7 @@ def dim_merger(grad, max_precond_dim, split: bool = False):
 
 
 def beta_debias(beta, step):
-    return 1 - (1 - beta) / (1 - beta ** step)
+    return 1 - (1 - beta) / (1 - beta**step)
 
 
 def eps_sqrt(item, eps):
@@ -164,8 +188,9 @@ 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]]):
+def _compilable_exp_avg_sq_(
+    state: List[Tensor], grad: List[Tensor], beta2: Tensor, eps: Tensor, out: List[Optional[Tensor]]
+):
     g32 = promote(grad)
     s32 = _lerp(state, torch._foreach_mul(g32, g32), beta2)
 
@@ -226,8 +251,9 @@ def _compilable_agc_(parameters: List[Tensor], gradients: List[Tensor], clip_val
     copy_stochastic_list_(gradients, g32)
 
 
-def adaptive_gradient_clipping_(parameters: List[Tensor], gradients: List[Tensor], clip_val: float,
-                                minimum: float = 1e-3, eps: float = 1e-8):
+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)
@@ -253,23 +279,24 @@ 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 +304,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 +317,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 +378,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 +430,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 +442,20 @@ 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}'
-    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])
+    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]
+    )
     copy_stochastic_(exp_avg, exp_avg_new)
 
     for q, q_new in zip(Q, new_qs):
@@ -453,11 +481,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 +517,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.")
 
 
@@ -554,6 +582,20 @@ def stochastic_add_(x: List[Tensor], y: List[Tensor], alpha: Union[float, int, T
     _compilable_stochastic_add_(x, y, alpha)
 
 
+@decorator_knowngood
+def _compilable_stochastic_add_divide_(x: List[Tensor], y: List[Tensor], alpha: Tensor, divisor: Tensor):
+    for x_, y_ in zip(x, y):
+        x32 = promote(x_)
+        y32 = promote(y_)
+        copy_stochastic_(x_, (x32 + y32 * alpha) / divisor)
+
+
+def stochastic_add_divide_(x: List[Tensor], y: List[Tensor], alpha: Union[float, int, Tensor] = 1, divisor: float = 1):
+    x, y = list_guard(x, y)
+    alpha, divisor = scalar_guard(alpha, divisor, x[0])
+    _compilable_stochastic_add_divide_(x, y, alpha, divisor)
+
+
 @decorator_knowngood
 def _compilable_stochastic_multiply_(x: List[Tensor], y: List[Tensor]):
     for x_, y_ in zip(x, y):
@@ -580,9 +622,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)
 
 
@@ -602,6 +644,20 @@ def promote(x):
     return x
 
 
+def promote_detach(x, should_promote):
+    if x is None:
+        return x
+    if should_promote:
+        x = promote(x)
+    return x.detach()
+
+
+def detach(x):
+    if isinstance(x, Tensor):
+        return x.detach()
+    return x
+
+
 def min_dtype(xs: List[Tensor]):
     dtypes = [x.dtype for x in xs]
     for d in (torch.float32, torch.bfloat16, torch.float16):
@@ -623,19 +679,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,34 +702,45 @@ 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
 
 
-def modify_closure(closure):
-    """
-    Modifies the closure function to use create_graph=True in backward().
+@contextlib.contextmanager
+def patch_backward():
+    @contextlib.contextmanager
+    def _inner(module):
+        original = module.backward
 
-    Args:
-        closure: The closure function passed to the optimizer.
+        signature = inspect.signature(original)
+
+        def patched_backward(*args, **kwargs):
+            new_kwargs = signature.bind(*args)
+            new_kwargs.apply_defaults()
+            new_kwargs = new_kwargs.arguments
+            new_kwargs.update(kwargs)
+            new_kwargs["create_graph"] = True
+            return original(**new_kwargs)
+
+        module.backward = patched_backward
+        yield
+        module.backward = original
+
+    with _inner(torch.Tensor), _inner(torch.autograd):
+        yield
 
-    Returns:
-        The return value of the modified closure.
-    """
 
-    def patched_backward(self, *args, **kwargs):
-        kwargs['create_graph'] = True
-        return original_backward(self, *args, **kwargs)
+def hasattr_none(obj, name):
+    return getattr(obj, name, None) is not None
 
-    original_backward = torch.Tensor.backward
 
-    with patch.object(torch.Tensor, 'backward', patched_backward):
-        return closure()
+class ExactHVPFailed(ValueError):
+    pass
 
 
 class StatefulOptimizer(torch.optim.Optimizer):
@@ -683,18 +750,22 @@ 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
     precond_schedule: Union[Callable, float, None] = None
     stochastic_schedule: bool = False
     finite_differences: bool = False
+    fallback_to_finite_differences: bool = True
+    _fallback_enabled: bool = False
+    hvp_interval: int = 1  # grad is faster initially, hvp later
 
     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,34 +781,51 @@ 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):
+    def split_p_and_g_in_group(
+        self,
+        group: dict,
+        skip_none: bool = True,
+        should_promote: bool = True,
+        beta1: float = -1.0,
+        raw: bool = False,
+    ):
         for p in group["params"]:
+            grad = getattr(p, "grad", None)
+            if grad is None and skip_none:
+                continue
+
+            p.grad = None
+
+            if raw:
+                yield p, grad
+                continue
+
             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)
-            p.grad = None
-
-            if grad is None:
-                grad = [getattr(pv, 'grad', None) for pv in p_views]
-            else:
-                grad = merge_group(group, grad)
-
-            for pv, g in zip(p_views, grad):
-                if skip_none and g is None:
-                    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)
+            vector = getattr(p, "vector", None)
+            hessian_vector = getattr(p, "hessian_vector", None)
+            p.vector = None
+            p.hessian_vector = None
+
+            grad, vs, hvs = [
+                [None] * len(p_views) if x is None else merge_group(group, x)  #
+                for x in (grad, vector, hessian_vector)
+            ]
+
+            for pv, g, v, hv in zip(p_views, grad, vs, hvs):
+                g = promote_detach(g, should_promote)
+                if beta1 >= 0 and group.get("mars", False):
+                    self.mars_correct_list(group, [pv], [g], group["mars_gamma"], beta1)
+                pv.vector = promote_detach(v, should_promote)
+                pv.hessian_vector = promote_detach(hv, should_promote)
                 yield pv, g
 
     def state_size(self) -> int:
@@ -759,48 +847,108 @@ 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 _finite_differences_hvp(self, closure):
+        with torch.enable_grad():
+            loss = closure()  # closure without retain_graph=True
+
+        grads = []
+        for group in self.param_groups:
+            for p, g in self.split_p_and_g_in_group(group, skip_none=True, raw=True):
+                grads.append(g)
+                p.vector = torch.randn_like(p)
+                p.orig = p.data.clone()
+                # 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)
+
+        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, raw=True):
+                p.grad = grads.pop(0)
+                stochastic_add_(g, p.grad, -1)  # technically, we have to divide by the scale here
+                p.hessian_vector = g
+                p.data.copy_(p.orig)
+                del p.orig
+        return loss
+
+    def _double_backward_hvp(self, closure):
+        with torch.enable_grad(), patch_backward():
+            loss = closure()
+
+        params, grads = [], []
+        for group in self.param_groups:
+            for p, g in self.split_p_and_g_in_group(group, skip_none=True, raw=True):
+                params.append(p)
+                grads.append(g)
+
+        if not params:
+            raise ValueError("No parameter has gradients")
+
+        vs = [torch.randn_like(p) for p in params]
+        with torch.enable_grad():
+            try:
+                hvs = torch.autograd.grad(grads, params, vs, create_graph=False, retain_graph=False, allow_unused=True)
+            except RuntimeError as e:
+                raise ExactHVPFailed(str(e.args))
+
+        unused = []
+        for p, g, v, hv in zip(params, grads, vs, hvs):
+            p.hessian_vector = detach(hv)
+            p.grad = detach(g)
+            p.vector = detach(v)
+            if hv is None:
+                unused.append(list(p.shape))
+
+        if unused:
+            raise ExactHVPFailed(f"Parameters with the following shapes have no 2nd order derivative: {unused}")
+
+        return loss
+
     def _handle_closure(self, closure):
         hessian_approx = self.hessian_approx and self._is_preconditioning
 
@@ -809,53 +957,41 @@ class StatefulOptimizer(torch.optim.Optimizer):
                 raise ValueError("Hessian approximation requires a closure.")
             return None
 
-        if not hessian_approx:
+        step = self._inner_group["total_hvp_steps"] = self._inner_group.get("total_hvp_steps", 0) + 1
+        if not hessian_approx or step % self.hvp_interval == 0:
             with torch.enable_grad():
                 loss = closure()
             return loss
 
-        if self.finite_differences:
-            with torch.enable_grad():
-                loss = closure()  # closure without retain_graph=True
-
-            grads = []
-            for group in self.param_groups:
-                for p, g in self.split_p_and_g_in_group(group, skip_none=True, should_promote=False):
-                    grads.append(g)
-                    p.vector = torch.randn_like(p)
-                    p.orig = p.data.clone()
-                    stochastic_add_(p.data, p.vector, tiny_bf16)
-        else:
-            with torch.enable_grad():
-                loss = modify_closure(closure)
-
-        if self.finite_differences:
-            with torch.enable_grad():
-                closure()
-
-            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)
-                    stochastic_add_(g, p.grad, -1)
-                    p.hessian_vector = g
-                    p.data.copy_(p.orig)
-                    del p.orig
-        else:
-            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 = 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)])
-            vs = [torch.randn_like(p) for p in params]
-            with torch.enable_grad():
-                hvs = torch.autograd.grad(grads, params, vs)
-
-            for p, g, v, hv in zip(params, grads, vs, hvs):
-                p.hessian_vector = hv
-                p.grad = g
-                p.vector = v
-
-        return loss
+        if self.finite_differences or self._fallback_enabled:
+            return self._finite_differences_hvp(closure)
+
+        try:
+            return self._double_backward_hvp(closure)
+        except NotImplementedError as e:
+            if not self.fallback_to_finite_differences:
+                raise
+            if not any(isinstance(arg, str) and _cudnn_double_backward_pattern.match(arg) for arg in e.args):
+                raise
+            warn_once(
+                "CUDNN doesn't support double-backward for some models (including RNNs). "  #
+                f"Falling back to finite_differences.\n{_fd_error}{e}"
+            )
+        except RuntimeError as e:
+            if not self.fallback_to_finite_differences:
+                raise
+            if not any(isinstance(arg, str) and _torch_compile_double_backward_pattern.match(arg) for arg in e.args):
+                raise
+            warn_once(
+                f"torch.compile does not support double-backward. Disabling it may be beneficial, depending on "
+                f"the model.\n{_fd_error}{e}"
+            )
+        except ExactHVPFailed as e:
+            if not self.fallback_to_finite_differences:
+                raise
+            warn_once(f"Exact HVP calculation failed.\n{_fd_error}{e}")
+        self._fallback_enabled = True
+        return self._handle_closure(closure)
 
     def step(self, closure: Optional[Callable] = None):
         if self.precond_schedule is None:
@@ -867,11 +1003,15 @@ 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()
-
+                for real, views in self.mapping.items():
+                    for tensor in (real, *views):
+                        for key in ("grad", "vector", "hessian_vector", "orig"):
+                            if hasattr(tensor, key):
+                                setattr(tensor, key, None)
         return loss
 
 
@@ -891,8 +1031,15 @@ 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):
+def _compilable_adam_(
+    exp_avg: List[Tensor],
+    exp_avg_sq: List[Tensor],
+    grad: List[Tensor],
+    beta1: Tensor,
+    beta2: Tensor,
+    step: Tensor,
+    eps: Tensor,
+):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -903,8 +1050,15 @@ def _compilable_adam_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: Lis
     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,
-          eps: float = 1e-8):
+def adam_(
+    exp_avg: List[Tensor],
+    exp_avg_sq: List[Tensor],
+    grad: List[Tensor],
+    beta1: float,
+    beta2: float,
+    step: int,
+    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)
@@ -912,9 +1066,20 @@ 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):
+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,
+):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -925,17 +1090,35 @@ def _fused_compilable_adam_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq:
     _compilable_update_(y, u32, decay, lr, caution, g32)
 
 
-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):
+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,
+):
     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)
 
 
 @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):
+def _compilable_laprop_(
+    exp_avg: List[Tensor],
+    exp_avg_sq: List[Tensor],
+    grad: List[Tensor],
+    beta1: Tensor,
+    beta2: Tensor,
+    step: Tensor,
+    eps: Tensor,
+):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -946,8 +1129,15 @@ def _compilable_laprop_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: L
     copy_stochastic_list_(grad, gp32)
 
 
-def laprop_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], beta1: float, beta2: float, step: int,
-            eps: float = 1e-8):
+def laprop_(
+    exp_avg: List[Tensor],
+    exp_avg_sq: List[Tensor],
+    grad: List[Tensor],
+    beta1: float,
+    beta2: float,
+    step: int,
+    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)
@@ -955,9 +1145,20 @@ 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):
+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,
+):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -968,9 +1169,20 @@ def _fused_compilable_laprop_(y: List[Tensor], exp_avg: List[Tensor], exp_avg_sq
     _compilable_update_(y, u32, decay, lr, caution, gp32)
 
 
-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):
+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,
+):
     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 +1190,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 +1209,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)
@@ -1044,8 +1256,9 @@ 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]]):
+def _compilable_update_(
+    p: List[Tensor], u: List[Tensor], decay: Tensor, lr: Tensor, caution: bool, 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_)
@@ -1055,8 +1268,9 @@ def _compilable_update_(p: List[Tensor], u: List[Tensor], decay: Tensor, lr: Ten
         copy_stochastic_(p_, p32_)
 
 
-def update_param_(param: List[Tensor], update: List[Tensor], lr: float, decay: float, caution: bool = False,
-                  grad: List[Tensor] = None):
+def update_param_(
+    param: List[Tensor], update: List[Tensor], lr: float, decay: float, caution: bool = False, grad: List[Tensor] = None
+):
     param, update, grad = list_guard(param, update, grad)
     lr = scalar_guard(lr, param[0])
     if not caution:
@@ -1064,38 +1278,117 @@ 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 stable_exp(x: Tensor):
+    # fp16:
+    #   exp(x) is stable in [-17, 11]
+    #   `stable_exp` extends to [-17, 17]
+    #   average error (in [-10, 10]) increased from 2.288e-3 to 2.299e-3
+    # fp32:
+    #   exp(x) is stable in [-103, 88]
+    #   `stable_exp` extends to [-103, 103]
+    #   average error (in [-87, 87]) reduced from 3.309-06 to 3.224-06
+    return torch.where(x > 0, 1 / (-x).exp(), x.exp())
+
+
+@decorator_knowngood
+def mean_root(x: torch.Tensor, pow: float, eps=1e-12):
+    # 1 / (mean(x ** pow) ** (1 / pow / 2))
+    log_x = x.double().abs().clamp(min=eps).log()
+    log_mean_x_pow = (log_x * pow).logsumexp(dim=0) - math.log(x.numel())
+    return stable_exp(-log_mean_x_pow / pow / 2)
+
+
+@decorator_knowngood
+def divided_root(x: torch.Tensor, y: torch.Tensor, pow0: float, pow1: float, eps=1e-12):
+    # mean(x ** pow0) ** (1 / pow0 / 2) / mean(y ** pow1) ** (1 / pow1 / 2)
+    log_x = x.double().abs().clamp(min=eps).log()
+    log_y = y.double().abs().clamp(min=eps).log()
+
+    x_normed = (log_x * pow0).logsumexp(dim=0) - math.log(x.numel())
+    x_normed = x_normed / pow0 / 2
+
+    y_normed = (log_y * pow1).logsumexp(dim=0) - math.log(y.numel())
+    y_normed = y_normed / pow1 / 2
+
+    return stable_exp(x_normed - y_normed)
+
+
+def precond_init_scale(scale, scale_scale, grad, hessian_vector, vector, scale_max: float = 1e6):
+    automatic_scale = True
+    manual_hint = " Set it manually using `precond_init_scale=0.1`"
+    if scale is not None:
+        automatic_scale = False
+        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 and scale_scale != 1:
+            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."
+            )
+    elif hessian_vector is None:
+        scale = mean_root(grad, 4) * scale_scale
+    else:
+        scale = divided_root(vector, hessian_vector, 2, 4) * scale_scale
+    if isinstance(scale, torch.Tensor):
+        scale = scale.item()  # slow, but necessary
+    if np.isfinite(scale):
+        if scale > scale_max or scale < 1 / scale_max:
+            warn_once(f"The computed precond_init_scale {scale} is outside of the expected range.{manual_hint}")
+        return scale
+    if not automatic_scale:
+        raise ValueError("The manually set precond_init_scale is not finite")
+
+    for x in (grad, hessian_vector, vector):
+        if x is None:
+            continue
+        if torch.allclose(x, torch.zeros_like(x)).item():
+            raise ValueError(f"Grad or HVP is all 0s, causing NaNs in precond_init_scale computation.{manual_hint}")
+        if not torch.isfinite(x).all().item():
+            raise ValueError("Grad or HVP is not finite")
+    raise ValueError(f"Computed precond_init_scale is not finite.{manual_hint}")
+
+
+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 +1396,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))
 
@@ -1119,8 +1412,10 @@ def init_Q_exprs(t, scale, max_size, min_ndim_triangular, memory_save_mode, dtyp
     elif memory_save_mode == "all_diag":
         dim_diag = [True for _ in shape]
     else:
-        raise ValueError(f"Invalid memory_save_mode: {memory_save_mode}, must be one of "
-                         "[None, 'one_diag', 'all_diag', 'smart_one_diag']")
+        raise ValueError(
+            f"Invalid memory_save_mode: {memory_save_mode}, must be one of "
+            "[None, 'one_diag', 'all_diag', 'smart_one_diag']"
+        )
 
     Q = []
     piece1A, piece2A, piece3A = ([], "", "")
@@ -1129,7 +1424,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 +1438,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 +1453,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,37 +1465,207 @@ def psgd_balance_Q(Q_in):
     torch._foreach_mul_(Q_in, list(norms))
 
 
-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)
+@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:
-        conjB = V.permute(*range(1, order), 0).to(promote(G.dtype))
-    Q = [promote(q) for q in Q]
+        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
+
+
+@decorator_knowngood
+def flatten(x: List[Tensor], remaining: int = 0) -> Tensor:
+    last_dim = x[0].shape[-remaining:] if remaining else []
+    return torch.cat([i.reshape(-1, *last_dim) for i in x], 0)
+
+
+@decorator_knowngood
+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)
+
+
+@decorator_knowngood
+def casted_einsum(expr: str, *args: Tensor) -> Tensor:
+    md = min_dtype(args)
+    return torch.einsum(expr, *[a.to(md) for a in args]).to(args[-1].dtype)
+
+
+def psgd_calc_A_and_conjB(exprA, G, Q, conjB):  # conjB ("V", "vector") == randn during hvp/whitening
+    order = G.dim()
+    if order > 1:
+        conjB = conjB.view_as(G).permute(*range(1, order), 0)
+    conjB = conjB.to(promote(G.dtype))
+    A = casted_einsum(exprA, *Q, G)
     for i, q in enumerate(Q):
+        q = promote(q)
         if q.dim() <= 1:
             conjB /= q
         else:
-            conjB = torch.linalg.solve_triangular(q, conjB.reshape(-1, q.size(0)), upper=True, left=False).reshape_as(
-                conjB)
+            solved = torch.linalg.solve_triangular(q, conjB.reshape(-1, q.size(0)).contiguous(), upper=True, left=False)
+            conjB = solved.reshape_as(conjB)
         if i < order - 1:
-            conjB = torch.transpose(conjB, i, order - 1)
+            conjB = conjB.transpose(i, -1)
     return A, conjB
 
 
 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 +1682,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:
@@ -1225,9 +1690,12 @@ def psgd_update_precond(Q, exprs, G, precond_lr, oq, store_triu_as_line, V):
             term1 /= torch.where(norm > 0, psgd_lb(term2, norm), norm).clamp_(tiny_bf16)
             term1 = torch.mm(term1, q.to(term1.dtype))
         if store_triu_as_line:
-            term1 = triu_to_line([term1])[0][1]
-            o = o[1]
-        stochastic_add_(o, term1, -1)
+            term1 = triu_to_line([term1])[0][1]  # Convert update to line format
+            # Apply update directly to the tensor part of the state tuple o[1]
+            stochastic_add_(o[1], term1, -1)
+        else:
+            # Apply update to the state tensor o
+            stochastic_add_(o, term1, -1)
 
 
 @decorator_knowngood
@@ -1245,7 +1713,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)
 
@@ -1437,12 +1905,13 @@ def warn_once(msg):
         _warned.add(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
+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']:
+        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
@@ -1450,8 +1919,9 @@ 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):
+def precond_grad_cached_(
+    expr: str, ea: Tensor, *cached_q: Tensor, caution: bool = False, grad: Optional[Tensor] = None, cast: bool = True
+):
     if caution:
         ea = _compilable_cautioning(grad, ea)
     md = min_dtype(list(cached_q) + [ea])
@@ -1564,18 +2034,86 @@ 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
 
 
+@decorator_knowngood
+def _compilable_d_adapt_(grads: List[Tensor], update: List[Tensor], state: List[Tensor], delta: List[Tensor]):
+    for g_, u_, s_, d_ in zip(grads, update, state, delta):
+        g, u, s, d = promote(g_), promote(u_), promote(s_), promote(d_)
+        next_d = d * (g * s).sum()
+        s = s + u * d
+        next_d = next_d / s.abs().sum()
+        next_d = torch.maximum(next_d, d)
+        copy_stochastic_(u_, u * d)
+        copy_stochastic_(d_, next_d)
+        copy_stochastic_(s_, s)
+
+
+def d_adaptation(grads: List[Tensor], update: List[Tensor], state: List[Tensor], delta: List[Tensor]):
+    grads, update, state, delta = list_guard(grads, update, state, delta)
+    _compilable_d_adapt_(grads, update, state, delta)
+
+
+@decorator_knowngood
+def _compilable_lr_adapt_(
+    grads: List[Tensor], update: List[Tensor], state: List[Tensor], delta: List[Tensor], lr_lr: Tensor
+):
+    for g_, u_, s_, d_ in zip(grads, update, state, delta):
+        g, u, s, d = promote(g_), promote(u_), promote(s_), promote(d_)
+        lr_grad = d.sigmoid()
+        lr_grad = lr_grad * (1 - lr_grad)
+        lr_grad = lr_grad * (s * g).mean()
+        d = d - lr_grad * lr_lr
+        copy_stochastic_(d_, d)
+        copy_stochastic_(u_, u * d.sigmoid())
+        copy_stochastic_(s_, u)
+
+
+def lr_adaptation(grads: List[Tensor], update: List[Tensor], state: List[Tensor], delta: List[Tensor], lr_lr: float):
+    grads, update, state, delta = list_guard(grads, update, state, delta)
+    lr_lr = scalar_guard(lr_lr, grads[0])
+    _compilable_lr_adapt_(grads, update, state, delta, lr_lr)
+
+
+@decorator_knowngood
+def _compilable_pointwise_lr_adapt_(
+    grads: List[Tensor], update: List[Tensor], state: List[Tensor], delta: List[Tensor], lr_lr: Tensor
+):
+    for g_, u_, s_, d_ in zip(grads, update, state, delta):
+        g, u, s, d = promote(g_), promote(u_), promote(s_), promote(d_)
+        lr_grad = d.sigmoid()
+        lr_grad = lr_grad * (1 - lr_grad)
+        lr_grad = lr_grad * s * g
+        d = d - lr_grad * lr_lr
+        copy_stochastic_(d_, d)
+        copy_stochastic_(u_, u * d.sigmoid())
+        copy_stochastic_(s_, u)
+
+
+def pointwise_lr_adaptation(
+    grads: List[Tensor], update: List[Tensor], state: List[Tensor], delta: List[Tensor], lr_lr: float
+):
+    grads, update, state, delta = list_guard(grads, update, state, delta)
+    lr_lr = scalar_guard(lr_lr, grads[0])
+    _compilable_lr_adapt_(grads, update, state, delta, lr_lr)
+
+
 def hook_optimizer_into_model(model, optimizer, *args, **kwargs):
     optimizers = {}
 
@@ -1598,8 +2136,9 @@ 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.")
+    o.step = functools.partial(
+        warn_once, msg="You're trying to call `step` on a fused optimizer. This will not do anything."
+    )
 
     def _step(p: Tensor):
         seen_params.add(p)