torchzero 0.3.10__py3-none-any.whl → 0.3.11__py3-none-any.whl

torchzero/modules/optimizers/soap.py

@@ -24,11 +24,9 @@ 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:
+        if mat is not None and 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)
 
@@ -40,8 +38,7 @@ 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:
+        if mat is not None and len(mat) > 0:
             tensors = torch.tensordot(tensors, mat,dims=[[0], [1]]) # pyright:ignore[reportArgumentType]
         else:
             permute_order = list(range(1, len(tensors.shape))) + [0]
@@ -59,8 +56,7 @@ def get_orthogonal_matrix(mat: list[torch.Tensor | None]):
     float_data = False
     original_type = original_device = None
     for m in mat:
-        if m is None: continue
-        if len(m) == 0:
+        if m is None or len(m) == 0:
             matrix.append([])
             continue
         if m.dtype != torch.float:
@@ -100,13 +96,11 @@ def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | N
     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:
+        if m is None or len(m) == 0:
             matrix.append([])
             orth_matrix.append([])
             continue
+        assert o is not None
         if m.data.dtype != torch.float:
             original_type = m.data.dtype
             original_device = m.data.device
@@ -156,6 +150,24 @@ class SOAP(Transform):
             learning rate. Defaults to 1.
         bias_correction (bool, optional):
             enables adam bias correction. Defaults to True.
+
+    Examples:
+        SOAP:
+
+        .. code-block:: python
+
+            opt = tz.Modular(model.parameters(), tz.m.SOAP(), tz.m.LR(1e-3))
+
+        Stabilized SOAP:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.SOAP(),
+                tz.m.NormalizeByEMA(max_ema_growth=1.2),
+                tz.m.LR(1e-2)
+            )
     """
     def __init__(
         self,
@@ -187,7 +199,7 @@ class SOAP(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         updates = []
         # update preconditioners
         for i,(p,t, state, setting) in enumerate(zip(params, tensors, states, settings)):
@@ -200,7 +212,7 @@ class SOAP(Transform):
             # 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)
+                state["exp_avg_sq_projected"] = torch.zeros_like(t)
 
                 if not precondition_1d and t.ndim <= 1:
                     state['GG'] = []
@@ -230,22 +242,20 @@ class SOAP(Transform):
             # 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"]
+            exp_avg_sq_projected: torch.Tensor = state["exp_avg_sq_projected"]
 
             exp_avg.lerp_(t, 1-beta1)
 
             if t_projected is None:
-                exp_avg_sq.mul_(beta2).addcmul_(t, t, value=1-beta2)
+                exp_avg_sq_projected.mul_(beta2).addcmul_(t, t, value=1-beta2)
             else:
-                exp_avg_sq.mul_(beta2).addcmul_(t_projected, t_projected, value=1-beta2)
+                exp_avg_sq_projected.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 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 = }')
 
@@ -273,6 +283,6 @@ class SOAP(Transform):
             if state['GG'] is not None:
                 update_soap_covariances_(t, state['GG'], shampoo_beta)
                 if state['step'] % setting['precond_freq'] == 0:
-                    state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, state['GG'], state['Q'])
+                    state['Q'], state['exp_avg_sq_projected'] = get_orthogonal_matrix_QR(exp_avg_sq_projected, state['GG'], state['Q'])
 
         return updates

torchzero/modules/optimizers/sophia_h.py

@@ -35,6 +35,74 @@ def sophia_H(
 
 
 class SophiaH(Module):
+    """SophiaH optimizer from https://arxiv.org/abs/2305.14342
+
+    This is similar to Adam, but the second momentum is replaced by an exponential moving average of randomized hessian diagonal estimates, and the update is agressively clipped.
+
+    .. note::
+        In most cases SophiaH should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply SophiaH preconditioning to another module's output.
+
+    .. note::
+        If you are using gradient estimators or reformulations, set :code:`hvp_method` to "forward" or "central".
+
+    .. note::
+        This module requires the a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients for calculating HVPs.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    Args:
+        beta1 (float, optional): first momentum. Defaults to 0.96.
+        beta2 (float, optional): momentum for hessian diagonal estimate. Defaults to 0.99.
+        update_freq (int, optional):
+            frequency of updating hessian diagonal estimate via a hessian-vector product. Defaults to 10.
+        precond_scale (float, optional):
+            scale of the preconditioner. Defaults to 1.
+        clip (float, optional):
+            clips update to (-clip, clip). Defaults to 1.
+        eps (float, optional):
+            clips hessian diagonal esimate to be no less than this value. Defaults to 1e-12.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        h (float, optional): finite difference step size if :code:`hvp_method` is "forward" or "central". Defaults to 1e-3.
+        n_samples (int, optional):
+            number of hessian-vector products with random vectors to evaluate each time when updating
+            the preconditioner. Larger values may lead to better hessian diagonal estimate. Defaults to 1.
+        seed (int | None, optional): seed for random vectors. Defaults to None.
+        inner (Chainable | None, optional): preconditioning is applied to the output of this module. Defaults to None.
+
+    Examples:
+        Using SophiaH:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.SophiaH(),
+                tz.m.LR(0.1)
+            )
+
+        SophiaH preconditioner can be applied to any other module by passing it to the :code:`inner` argument.
+        Turn off SophiaH's first momentum to get just the preconditioning. Here is an example of applying
+        SophiaH preconditioning to nesterov momentum (:code:`tz.m.NAG`):
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.SophiaH(beta1=0, inner=tz.m.NAG(0.96)),
+                tz.m.LR(0.1)
+            )
+
+    """
     def __init__(
         self,
         beta1: float = 0.96,
@@ -85,23 +153,12 @@ class SophiaH(Module):
         h = None
         if step % update_freq == 0:
 
-            grad=None
+            rgrad=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 = var.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=var.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)
+                Hvp, rgrad = self.Hvp(u, at_x0=True, var=var, rgrad=rgrad, hvp_method=hvp_method,
+                                     h=fd_h, normalize=True, retain_grad=i < n_samples-1)
 
                 if h is None: h = Hvp
                 else: torch._foreach_add_(h, Hvp)

torchzero/modules/projections/__init__.py

@@ -1,5 +1,3 @@
-from .projection import Projection
-from .fft import FFTProjection
-from .structural import VectorProjection, TensorizeProjection, BlockPartition, TensorNormsProjection
-
+from .projection import ProjectionBase, VectorProjection, ScalarProjection
+from .cast import To, ViewAsReal
 # from .galore import GaLore

torchzero/modules/projections/cast.py

@@ -0,0 +1,51 @@
+import torch
+from .projection import ProjectionBase
+from ...core import Chainable
+
+class To(ProjectionBase):
+    """Cast modules to specified device and dtype"""
+    def __init__(self, modules: Chainable, dtype: torch.dtype | None, device:torch.types.Device | None = None):
+        defaults = dict(dtype=dtype, device=device)
+        super().__init__(modules, project_update=True, project_params=True, project_grad=True, defaults=defaults)
+
+    @torch.no_grad
+    def project(self, tensors, params, grads, loss, states, settings, current):
+        casted = []
+        for tensor, state, setting in zip(tensors,states, settings):
+            state['dtype'] = tensor.dtype
+            state['device'] = tensor.device
+            tensor = tensor.to(dtype=setting['dtype'], device=setting['device'])
+            casted.append(tensor)
+        return casted
+
+    @torch.no_grad
+    def unproject(self, projected_tensors, params, grads, loss, states, settings, current):
+        uncasted = []
+        for tensor, state in zip(projected_tensors, states):
+            tensor = tensor.to(dtype=state['dtype'], device=state['device'])
+            uncasted.append(tensor)
+        return uncasted
+
+
+class ViewAsReal(ProjectionBase):
+    """View complex tensors as real tensors. Doesn't affect tensors that are already."""
+    def __init__(self, modules: Chainable):
+        super().__init__(modules, project_update=True, project_params=True, project_grad=True, defaults=None)
+
+    @torch.no_grad
+    def project(self, tensors, params, grads, loss, states, settings, current):
+        views = []
+        for tensor, state in zip(tensors,states):
+            is_complex = torch.is_complex(tensor)
+            state['is_complex'] = is_complex
+            if is_complex: tensor = torch.view_as_real(tensor)
+            views.append(tensor)
+        return views
+
+    @torch.no_grad
+    def unproject(self, projected_tensors, params, grads, loss, states, settings, current):
+        un_views = []
+        for tensor, state in zip(projected_tensors, states):
+            if state['is_complex']: tensor = torch.view_as_complex(tensor)
+            un_views.append(tensor)
+        return un_views

torchzero/modules/projections/galore.py

@@ -7,4 +7,6 @@ from typing import Any, Literal
 import torch
 
 from ...core import Chainable, Module, Var
-from .projection import Projection
+from .projection import ProjectionBase
+
+# TODO