torchzero 0.1.7__py3-none-any.whl → 0.3.1__py3-none-any.whl

torchzero/modules/optimizers/soap.py

@@ -0,0 +1,286 @@
+from operator import itemgetter
+
+import torch
+
+from ...core import Chainable, Transform, apply
+from ...modules.optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
+
+@torch.no_grad
+def update_soap_covariances_(
+    grad: torch.Tensor,
+    GGs_: list[torch.Tensor | None],
+    beta: float | None,
+):
+    for i, GG in enumerate(GGs_):
+        if GG is None: continue
+
+        axes = list(range(i)) + list(range(i + 1, grad.ndim)) # this works fine with 1d params
+        if beta is None: GG.add_(torch.tensordot(grad, grad, (axes, axes))) # pyright:ignore[reportArgumentType]
+        else: GG.lerp_(torch.tensordot(grad, grad, (axes, axes)), 1-beta) # pyright:ignore[reportArgumentType]
+
+@torch.no_grad
+def project(tensors: torch.Tensor, Q: list[torch.Tensor | None]):
+    """
+    Projects the gradient to the eigenbases of the preconditioner.
+    """
+    for mat in Q:
+        if mat is None: continue
+        if len(mat) > 0:
+            tensors = torch.tensordot(tensors, mat, dims=[[0], [0]]) # pyright:ignore[reportArgumentType]
+        else:
+            # I don't understand this part but it is in https://github.com/nikhilvyas/SOAP/blob/main/soap.py
+            permute_order = list(range(1, len(tensors.shape))) + [0]
+            tensors = tensors.permute(permute_order)
+
+    return tensors
+
+@torch.no_grad
+def project_back(tensors: torch.Tensor, Q: list[torch.Tensor| None]):
+    """
+    Projects the gradient back to the original space.
+    """
+    for mat in Q:
+        if mat is None: continue
+        if len(mat) > 0:
+            tensors = torch.tensordot(tensors, mat,dims=[[0], [1]]) # pyright:ignore[reportArgumentType]
+        else:
+            permute_order = list(range(1, len(tensors.shape))) + [0]
+            tensors = tensors.permute(permute_order)
+
+    return tensors
+
+# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py
+@torch.no_grad
+def get_orthogonal_matrix(mat: list[torch.Tensor | None]):
+    """
+    Computes the eigenbases of the preconditioner using torch.linalg.eigh decomposition.
+    """
+    matrix = []
+    float_data = False
+    original_type = original_device = None
+    for m in mat:
+        if m is None: continue
+        if len(m) == 0:
+            matrix.append([])
+            continue
+        if m.dtype != torch.float:
+            original_type = m.dtype
+            original_device = m.device
+            matrix.append(m.float())
+        else:
+            float_data = True
+            matrix.append(m)
+
+    final = []
+    for m in matrix:
+        if len(m) == 0:
+            final.append([])
+            continue
+        try:
+            _, Q = torch.linalg.eigh(m+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
+        except Exception:
+            _, Q = torch.linalg.eigh(m.to(torch.float64)+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
+            Q = Q.to(m.dtype)
+        Q = torch.flip(Q, [1])
+
+        if not float_data:
+            Q = Q.to(original_device).type(original_type)
+        final.append(Q)
+    return final
+
+# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py#L240
+@torch.no_grad
+def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | None], Q_list: list[torch.Tensor | None]):
+    """
+    Computes the eigenbases of the preconditioner using one round of power iteration
+    followed by torch.linalg.qr decomposition.
+    """
+    matrix = []
+    orth_matrix = []
+    float_data = False
+    original_type = original_device = None
+    for m,o in zip(GG, Q_list):
+        if m is None: continue
+        assert o is not None
+
+        if len(m) == 0:
+            matrix.append([])
+            orth_matrix.append([])
+            continue
+        if m.data.dtype != torch.float:
+            original_type = m.data.dtype
+            original_device = m.data.device
+            matrix.append(m.data.float())
+            orth_matrix.append(o.data.float())
+        else:
+            float_data = True
+            matrix.append(m.data.float())
+            orth_matrix.append(o.data.float())
+
+    final = []
+    for ind, (m,o) in enumerate(zip(matrix, orth_matrix)):
+        if len(m)==0:
+            final.append([])
+            continue
+        est_eig = torch.diag(o.T @ m @ o)
+        sort_idx = torch.argsort(est_eig, descending=True)
+        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
+        o = o[:,sort_idx]
+        power_iter = m @ o
+        Q, _ = torch.linalg.qr(power_iter) # pylint:disable=not-callable
+
+        if not float_data:
+            Q = Q.to(original_device).type(original_type)
+        final.append(Q)
+
+    return final, exp_avg_sq
+
+class SOAP(Transform):
+    """SOAP (ShampoO with Adam in the Preconditioner's eigenbasis from https://arxiv.org/abs/2409.11321).
+
+    Args:
+        beta1 (float, optional): beta for first momentum. Defaults to 0.95.
+        beta2 (float, optional): beta for second momentum. Defaults to 0.95.
+        shampoo_beta (float | None, optional):
+            beta for covariance matrices accumulators. Can be None, then it just sums them like Adagrad (which works worse). Defaults to 0.95.
+        precond_freq (int, optional): How often to update the preconditioner. Defaults to 10.
+        merge_small (bool, optional): Whether to merge small dims. Defaults to True.
+        max_dim (int, optional): Won't precondition dims larger than this. Defaults to 2_000.
+        precondition_1d (bool, optional):
+            Whether to precondition 1d params (SOAP paper sets this to False). Defaults to True.
+        eps (float, optional):
+            epsilon for dividing first momentum by second. Defaults to 1e-8.
+        decay (float | None, optional):
+            Decays covariance matrix accumulators, this may be useful if `shampoo_beta` is None. Defaults to None.
+        unprojected_exp_avg (bool, optional):
+            whether to update first momentum in unprojected space. Both true and false work and lead to different
+            results but True usually works better. Defaults to True.
+        bias_correction (bool, optional):
+            enables adam bias correction. Defaults to True.
+    """
+    def __init__(
+        self,
+        beta1: float = 0.95,
+        beta2: float = 0.95,
+        shampoo_beta: float | None = 0.95,
+        precond_freq: int = 10,
+        merge_small: bool = True,
+        max_dim: int = 2_000,
+        precondition_1d: bool = True,
+        eps: float = 1e-8,
+        decay: float | None = None,
+        alpha: float = 1,
+        unprojected_exp_avg: bool = True,
+        bias_correction: bool = True,
+    ):
+        defaults = dict(
+            beta1=beta1,
+            beta2=beta2,
+            shampoo_beta=shampoo_beta,
+            precond_freq=precond_freq,
+            merge_small=merge_small,
+            max_dim=max_dim,
+            precondition_1d=precondition_1d,
+            eps=eps,
+            decay=decay,
+            unprojected_exp_avg=unprojected_exp_avg,
+            bias_correction=bias_correction,
+            alpha=alpha,
+        )
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        updates = []
+        # update preconditioners
+        for i,(p,t) in enumerate(zip(params, tensors)):
+            state = self.state[p]
+            settings = self.settings[p]
+            beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps, unprojected_exp_avg,alpha = itemgetter(
+                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'unprojected_exp_avg','alpha')(settings)
+
+            if merge_small:
+                t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
+
+            # initialize state on 1st step
+            if 'GG' not in state:
+                state["exp_avg"] = torch.zeros_like(t)
+                state["exp_avg_sq"] = torch.zeros_like(t)
+
+                if not precondition_1d and t.ndim <= 1:
+                    state['GG'] = []
+
+                else:
+                    state['GG'] = [torch.zeros(s, s, dtype=t.dtype, device=t.device) if 1<s<max_dim else None for s in t.shape]
+
+                # either scalar parameter, 1d with precondition_1d=False, or all dims are too big.
+                if len([i is not None for i in state['GG']]) == 0:
+                    state['GG'] = None
+
+                if state['GG'] is not None:
+                    update_soap_covariances_(t, GGs_=state['GG'], beta=shampoo_beta)
+                    state['Q'] = get_orthogonal_matrix(state['GG'])
+
+                state['step'] = 0
+                updates.append(tensors[i].sign().div_(10))
+                # updates.append(tensors[i] / tensors[i].abs().sum())
+                continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
+                # I use scaled update instead as to not mess up with next modules.
+
+            # Projecting gradients to the eigenbases of Shampoo's preconditioner
+            # i.e. projecting to the eigenbases of matrices in state['GG']
+            t_projected = None
+            if state['GG'] is not None:
+                t_projected = project(t, state['Q'])
+
+            # exponential moving averages
+            # this part could be foreached but I will do that at some point its not a big difference compared to preconditioning
+            exp_avg: torch.Tensor = state["exp_avg"]
+            exp_avg_sq: torch.Tensor = state["exp_avg_sq"]
+
+            if unprojected_exp_avg or t_projected is None:
+                exp_avg.lerp_(t, 1-beta1)
+            else:
+                exp_avg.lerp_(t_projected, 1-beta1)
+
+            if t_projected is None:
+                exp_avg_sq.mul_(beta2).addcmul_(t, t, value=1-beta2)
+            else:
+                exp_avg_sq.mul_(beta2).addcmul_(t_projected, t_projected, value=1-beta2)
+
+            # project exponential moving averages if they are accumulated unprojected
+            exp_avg_projected = exp_avg
+            if unprojected_exp_avg and t_projected is not None:
+                exp_avg_projected = project(exp_avg, state['Q'])
+
+            exp_avg_sq_projected = exp_avg_sq
+
+            denom = exp_avg_sq_projected.sqrt().add_(eps)
+            # print(f'{t_projected = }, {exp_avg = }, {exp_avg_projected = }, {exp_avg_sq = }, {exp_avg_sq_projected = }, {denom = }')
+
+            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
+            # to the original space
+            update = exp_avg_projected / denom
+            if t_projected is not None:
+                update = project_back(update, state["Q"])
+
+            if settings['bias_correction']:
+                bias_correction1 = 1.0 - beta1 ** (state["step"]+1)
+                bias_correction2 = 1.0 - beta2 ** (state["step"]+1)
+                update *= ((bias_correction2 ** .5) / bias_correction1) * alpha
+            elif alpha is not None:
+                update *= alpha
+
+            if merge_small:
+                update = _unmerge_small_dims(update, state['flat_sizes'], state['sort_idxs'])
+
+            updates.append(update)
+            state["step"] += 1
+
+            # Update is done after the gradient step to avoid using current gradients in the projection.
+            if state['GG'] is not None:
+                update_soap_covariances_(t, state['GG'], shampoo_beta)
+                if state['step'] % settings['precond_freq'] == 0:
+                    state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, state['GG'], state['Q'])
+
+        return updates

torchzero/modules/optimizers/sophia_h.py

@@ -0,0 +1,129 @@
+from typing import Literal
+from collections.abc import Callable
+import torch
+
+from ...core import Module, Target, Transform, Chainable, apply
+from ...utils import NumberList, TensorList, as_tensorlist
+from ...utils.derivatives import hvp, hvp_fd_forward, hvp_fd_central
+
+def sophia_H(
+    tensors: TensorList,
+    h: TensorList | None,
+    exp_avg_: TensorList,
+    h_exp_avg_: TensorList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    update_freq: int,
+    precond_scale: float | NumberList,
+    clip: float | NumberList,
+    eps: float | NumberList,
+    step: int
+):
+    # momentum
+    exp_avg_.lerp_(tensors, 1-beta1)
+
+    # update preconditioner
+    if step % update_freq == 0:
+        assert h is not None
+        h_exp_avg_.lerp_(h, 1-beta2)
+
+    else:
+        assert h is None
+
+    denom = (h_exp_avg_ * precond_scale).clip_(min=eps)
+    return (exp_avg_ / denom).clip_(-clip, clip)
+
+
+class SophiaH(Module):
+    def __init__(
+        self,
+        beta1: float = 0.96,
+        beta2: float = 0.99,
+        update_freq: int = 10,
+        precond_scale: float = 1,
+        clip: float = 1,
+        eps: float = 1e-12,
+        hvp_method: Literal['autograd', 'forward', 'central'] = 'autograd',
+        fd_h: float = 1e-3,
+        n_samples = 1,
+        seed: int | None = None,
+        inner: Chainable | None = None
+    ):
+        defaults = dict(beta1=beta1, beta2=beta2, update_freq=update_freq, precond_scale=precond_scale, clip=clip, eps=eps, hvp_method=hvp_method, n_samples=n_samples, fd_h=fd_h, seed=seed)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = vars.params
+        settings = self.settings[params[0]]
+        hvp_method = settings['hvp_method']
+        fd_h = settings['fd_h']
+        update_freq = settings['update_freq']
+        n_samples = settings['n_samples']
+
+        seed = settings['seed']
+        generator = None
+        if seed is not None:
+            if 'generator' not in self.global_state:
+                self.global_state['generator'] = torch.Generator(params[0].device).manual_seed(seed)
+            generator = self.global_state['generator']
+
+        beta1, beta2, precond_scale, clip, eps = self.get_settings(
+            'beta1', 'beta2', 'precond_scale', 'clip', 'eps', params=params, cls=NumberList)
+
+        exp_avg, h_exp_avg = self.get_state('exp_avg', 'h_exp_avg', params=params, cls=TensorList)
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        closure = vars.closure
+        assert closure is not None
+
+        h = None
+        if step % update_freq == 0:
+
+            grad=None
+            for i in range(n_samples):
+                u = [torch.randn(p.shape, device=p.device, dtype=p.dtype, generator=generator) for p in params]
+
+                if hvp_method == 'autograd':
+                    if grad is None: grad = vars.get_grad(create_graph=True)
+                    assert grad is not None
+                    Hvp = hvp(params, grad, u, retain_graph=i < n_samples-1)
+
+                elif hvp_method == 'forward':
+                    loss, Hvp = hvp_fd_forward(closure, params, u, h=fd_h, g_0=vars.get_grad(), normalize=True)
+
+                elif hvp_method == 'central':
+                    loss, Hvp = hvp_fd_central(closure, params, u, h=fd_h, normalize=True)
+
+                else:
+                    raise ValueError(hvp_method)
+
+                if h is None: h = Hvp
+                else: torch._foreach_add_(h, Hvp)
+
+            assert h is not None
+            if n_samples > 1: torch._foreach_div_(h, n_samples)
+
+        update = vars.get_update()
+        if 'inner' in self.children:
+            update = apply(self.children['inner'], tensors=update, params=params, grads=vars.grad, vars=vars)
+
+        vars.update = sophia_H(
+            tensors=TensorList(update),
+            h=TensorList(h) if h is not None else None,
+            exp_avg_=exp_avg,
+            h_exp_avg_=h_exp_avg,
+            beta1=beta1,
+            beta2=beta2,
+            update_freq=update_freq,
+            precond_scale=precond_scale,
+            clip=clip,
+            eps=eps,
+            step=step,
+        )
+        return vars

torchzero/modules/projections/__init__.py

@@ -0,0 +1,5 @@
+from .projection import Projection
+from .fft import FFTProjection
+from .structural import VectorProjection, TensorizeProjection, BlockPartition, TensorNormsProjection
+
+# from .galore import GaLore

torchzero/modules/projections/dct.py

@@ -0,0 +1,73 @@
+from typing import Literal
+import torch
+import torch_dct
+from .projection import Projection
+from ...core import Chainable
+
+def reverse_dims(t:torch.Tensor):
+    return t.permute(*reversed(range(t.ndim)))
+
+class DCTProjection(Projection):
+    # norm description copied from pytorch docstring
+    """Project update into Discrete Cosine Transform space, requires `torch_dct` library.
+
+    Args:
+        modules (Chainable): modules that will optimize the projected update.
+        dims (1, 2 or 3, optional):
+            applies DCT to first 1,2 or 3 dims, defaults to 3.
+        norm (str, optional):
+            Normalization mode.
+            * None - no normalization
+            * "ortho" - normalize by 1/sqrt(n)
+    """
+
+    def __init__(
+        self,
+        modules: Chainable,
+        dims: Literal[1, 2, 3] = 3,
+        norm=None,
+        project_update=True,
+        project_params=False,
+        project_grad=False,
+    ):
+        defaults = dict(dims=dims, norm=norm)
+        super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad, defaults=defaults)
+
+    @torch.no_grad
+    def project(self, tensors, vars, current):
+        settings = self.settings[vars.params[0]]
+        dims = settings['dims']
+        norm = settings['norm']
+
+        projected = []
+        for u in tensors:
+            u = reverse_dims(u)
+            dim = min(u.ndim, dims)
+
+            if dim == 1: dct = torch_dct.dct(u, norm = norm)
+            elif dim == 2: dct = torch_dct.dct_2d(u, norm=norm)
+            elif dim == 3: dct = torch_dct.dct_3d(u, norm=norm)
+            else: raise ValueError(f"Unsupported number of dimensions {dim}")
+
+            projected.append(dct)
+
+        return projected
+
+    @torch.no_grad
+    def unproject(self, tensors, vars, current):
+        settings = self.settings[vars.params[0]]
+        dims = settings['dims']
+        norm = settings['norm']
+
+        unprojected = []
+        for u in tensors:
+            dim = min(u.ndim, dims)
+
+            if dim == 1: idct = torch_dct.idct(u, norm = norm)
+            elif dim == 2: idct = torch_dct.idct_2d(u, norm=norm)
+            elif dim == 3: idct = torch_dct.idct_3d(u, norm=norm)
+            else: raise ValueError(f"Unsupported number of dimensions {dim}")
+
+            unprojected.append(reverse_dims(idct))
+
+        return unprojected

torchzero/modules/projections/fft.py

@@ -0,0 +1,73 @@
+import torch
+
+from ...core import Chainable
+from ...utils import vec_to_tensors
+from .projection import Projection
+
+
+class FFTProjection(Projection):
+    # norm description copied from pytorch docstring
+    """Project update into Fourrier space of real-valued inputs.
+
+    Args:
+        modules (Chainable): modules that will optimize the projected update.
+        one_d (bool, optional):
+            * If True, uses 1d fft on parameters concatenated into a vector.
+            * If False, uses n-dimensional fft on each parameter (default).
+        norm (str, optional):
+            Normalization mode.
+
+            * "forward" - normalize by 1/n
+            * "backward" - no normalization
+            * "ortho" - normalize by 1/sqrt(n) (making the FFT orthonormal)
+
+            Calling the backward transform (:func:`~torch.fft.irfft`) with the same
+            normalization mode will apply an overall normalization of ``1/n`` between
+            the two transforms. This is required to make :func:`~torch.fft.irfft`
+            the exact inverse.
+
+            Default is "backward" (no normalization).
+
+            The actual torch.fft.rfft default is None, so I set it to None too. I guess None and "backward"
+            are the same.
+    """
+
+    def __init__(
+        self,
+        modules: Chainable,
+        one_d: bool = False,
+        norm=None,
+        project_update=True,
+        project_params=False,
+        project_grad=False,
+    ):
+        defaults = dict(one_d=one_d, norm=norm)
+        super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad, defaults=defaults)
+
+    @torch.no_grad
+    def project(self, tensors, vars, current):
+        settings = self.settings[vars.params[0]]
+        one_d = settings['one_d']
+        norm = settings['norm']
+
+        # 1d fft, concatenate all parameters into a vector and calculate fft
+        if one_d:
+            vec = torch.cat([t.view(-1) for t in tensors])
+            self.global_state['length'] = len(vec)
+            return [torch.view_as_real(torch.fft.rfft(vec, norm=norm))] # pylint:disable=not-callable
+
+        # multidimensional fft for each parameter
+        return [torch.view_as_real(torch.fft.rfftn(t, norm=norm)) if t.numel() > 1 else t for t in tensors] # pylint:disable=not-callable
+
+    @torch.no_grad
+    def unproject(self, tensors, vars, current):
+        settings = self.settings[vars.params[0]]
+        one_d = settings['one_d']
+        norm = settings['norm']
+
+        if one_d:
+            vec = torch.view_as_complex(tensors[0])
+            unprojected_vec = torch.fft.irfft(vec, n=self.global_state['length'], norm=norm) # pylint:disable=not-callable
+            return vec_to_tensors(unprojected_vec, reference=vars.params)
+
+        return [torch.fft.irfftn(torch.view_as_complex(t.contiguous()), s=p.shape, norm=norm) if t.numel() > 1 else t for t, p in zip(tensors, vars.params)] # pylint:disable=not-callable

torchzero/modules/projections/galore.py

@@ -0,0 +1,10 @@
+import importlib.util
+import warnings
+from collections.abc import Callable, Mapping
+from operator import itemgetter
+from typing import Any, Literal
+
+import torch
+
+from ...core import Chainable, Module, Vars
+from .projection import Projection