heavyball 1.7.0__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 Callable, List, Optional, Tuple, Union
-from unittest.mock import patch
 
 import numpy as np
 import torch
@@ -15,13 +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
+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
 tiny_bf16 = torch.finfo(torch.bfloat16).tiny
+_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):
@@ -58,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_))
@@ -74,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
 
@@ -149,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):
@@ -157,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)
 
@@ -219,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)
@@ -259,9 +292,11 @@ def set_torch(benchmark_limit: int = 32, einsum_strategy: str = "auto"):
 
     # 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
@@ -408,7 +443,7 @@ def get_orthogonal_matrix_QR(GG: List[Tensor], Q: List[Tensor], exp_avg: Optiona
 
     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()]
+    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:
@@ -418,8 +453,9 @@ def get_orthogonal_matrix_QR(GG: List[Tensor], Q: List[Tensor], exp_avg: Optiona
     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])
+    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):
@@ -546,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):
@@ -594,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):
@@ -647,25 +711,36 @@ def project(grad, Q, back: bool):
     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)
 
-    Returns:
-        The return value of the modified closure.
-    """
+        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)
 
-    def patched_backward(self, *args, **kwargs):
-        kwargs["create_graph"] = True
-        return original_backward(self, *args, **kwargs)
+        module.backward = patched_backward
+        yield
+        module.backward = original
 
-    original_backward = torch.Tensor.backward
+    with _inner(torch.Tensor), _inner(torch.autograd):
+        yield
 
-    with patch.object(torch.Tensor, "backward", patched_backward):
-        return closure()
+
+def hasattr_none(obj, name):
+    return getattr(obj, name, None) is not None
+
+
+class ExactHVPFailed(ValueError):
+    pass
 
 
 class StatefulOptimizer(torch.optim.Optimizer):
@@ -682,6 +757,9 @@ class StatefulOptimizer(torch.optim.Optimizer):
     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})
@@ -708,29 +786,46 @@ class StatefulOptimizer(torch.optim.Optimizer):
         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
+            vector = getattr(p, "vector", None)
+            hessian_vector = getattr(p, "hessian_vector", None)
+            p.vector = None
+            p.hessian_vector = None
 
-            if grad is None:
-                grad = [getattr(pv, "grad", None) for pv in p_views]
-            else:
-                grad = merge_group(group, grad)
+            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 in zip(p_views, grad):
-                if skip_none and g is None:
-                    continue
-                if should_promote:
-                    g = promote(g)
+            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:
@@ -794,6 +889,66 @@ class StatefulOptimizer(torch.optim.Optimizer):
                         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
 
@@ -802,56 +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()
-                    # 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)
-
-        if self.finite_differences:
-            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)
-                    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
+        if self.finite_differences or self._fallback_enabled:
+            return self._finite_differences_hvp(closure)
 
-        return loss
+        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,7 +1007,11 @@ class StatefulOptimizer(torch.optim.Optimizer):
                 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
 
 
@@ -887,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)
 
@@ -899,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)
@@ -908,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)
 
@@ -921,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)
 
@@ -942,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)
@@ -951,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)
 
@@ -964,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)
@@ -1040,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_)
@@ -1051,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:
@@ -1076,28 +1294,74 @@ def _max_idx(x: List[int]):
 
 
 @decorator_knowngood
-def mean_root(x: torch.Tensor, pow: float):
-    return stochastic_round_(x, x.float().pow(pow).mean().pow(-1 / pow / 2))
+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 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 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):
+
+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:
+            "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.")
+                "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 hessian_vector is None:
-        return mean_root(grad, 4) * scale_scale
-    return divided_root(vector, hessian_vector, 2, 4) * scale_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):
@@ -1108,8 +1372,9 @@ def init_lra(grad, scale, scale_scale, rank, hessian_vector, vector, dtype=None)
     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):
+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.
@@ -1147,8 +1412,10 @@ def init_Q_exprs(grad, scale, scale_scale, max_size, min_ndim_triangular, memory
     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 = ([], "", "")
@@ -1213,8 +1480,16 @@ def low_rank_mm(U: Tensor, V: Tensor, x: Tensor) -> Tensor:
     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, ):
+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
     """
@@ -1293,7 +1568,7 @@ def lra_precond(U, V, d, g):
 
 
 @decorator_knowngood
-def dampen_grad(g: Tensor, damp: float = 2 ** -13):
+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
@@ -1306,7 +1581,7 @@ def apply_lra_update(params: List[Tensor], update: Tensor, U: Tensor, V: Tensor,
     update = update.flatten()
     for p in params:
         size = p.numel()
-        copy_stochastic_(p, update[start: start + size].view_as(p))
+        copy_stochastic_(p, update[start : start + size].view_as(p))
         start += size
 
 
@@ -1316,7 +1591,7 @@ def apply_flat_update(params: List[Tensor], update: Tensor):
     update = update.flatten()
     for p in params:
         size = p.numel()
-        copy_stochastic_(p, update[start: start + size].view_as(p))
+        copy_stochastic_(p, update[start : start + size].view_as(p))
         start += size
 
 
@@ -1326,7 +1601,7 @@ def apply_flat_add(params: List[Tensor], update: Tensor, alpha: Tensor):
     update = update.flatten()
     for p in params:
         size = p.numel()
-        stochastic_add_([p], [update[start: start + size].view_as(p)], alpha)
+        stochastic_add_([p], [update[start : start + size].view_as(p)], alpha)
         start += size
 
 
@@ -1337,16 +1612,19 @@ def extract_from_flat_update(params: List[Tensor], update: Tensor):
     update = update.flatten()
     for p in params:
         size = p.numel()
-        outputs.append(update[start: start + size].view_as(p))
+        outputs.append(update[start : start + size].view_as(p))
         start += size
     return outputs
 
 
+@decorator_knowngood
 def flatten(x: List[Tensor], remaining: int = 0) -> Tensor:
-    return torch.cat([i.flatten(0, -1 - remaining) for i in x], 0)
+    last_dim = x[0].shape[-remaining:] if remaining else []
+    return torch.cat([i.reshape(-1, *last_dim) for i in x], 0)
 
 
-def dampen_multiple(g: List[Tensor], damp: float = 2 ** -13):
+@decorator_knowngood
+def dampen_multiple(g: List[Tensor], damp: float = 2**-13):
     vs = []
     gs = []
     for g_ in g:
@@ -1356,22 +1634,27 @@ def dampen_multiple(g: List[Tensor], damp: float = 2 ** -13):
     return flatten(vs), flatten(gs)
 
 
-def psgd_calc_A_and_conjB(exprA, G, Q, V=None):
+@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 V is None:
-        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]
+    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
 
 
@@ -1407,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
@@ -1619,8 +1905,9 @@ 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"):
+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"])
@@ -1632,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])
@@ -1753,12 +2041,79 @@ def merge_group(group, *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 = {}
 
@@ -1781,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)