torchzero 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

torchzero/modules/adaptive/muon.py

@@ -10,8 +10,10 @@ from ...linalg.orthogonalize import orthogonalize as _orthogonalize, Orthogonali
 def reverse_dims(t:torch.Tensor):
     return t.permute(*reversed(range(t.ndim)))
 
-def _is_at_least_2d(p: torch.Tensor):
-    if (p.ndim >= 2) and (p.size(0) > 1) and (p.size(1) > 1): return True
+def _is_at_least_2d(p: torch.Tensor, channel_first:bool):
+    if p.ndim < 2: return False
+    if channel_first and (p.size(0) > 1) and (p.size(1) > 1): return True
+    if (not channel_first) and (p.size(-2) > 1) and (p.size(-1) > 1): return True
     return False
 
 def _orthogonalize_format(
@@ -19,6 +21,7 @@ def _orthogonalize_format(
     method: OrthogonalizeMethod,
     channel_first: bool,
 ):
+    """orthogonalize either 1st two dims if channel first or last two otherwise"""
     if channel_first:
         return reverse_dims(_orthogonalize(reverse_dims(tensor), method=method))
 
@@ -69,7 +72,7 @@ def orthogonalize_grads_(
             are considered batch dimensions.
     """
     for p in params:
-        if (p.grad is not None) and _is_at_least_2d(p.grad):
+        if (p.grad is not None) and _is_at_least_2d(p.grad, channel_first=channel_first):
             X = _orthogonalize_format(p.grad, method=method, channel_first=channel_first)
             if dual_norm_correction: X = _dual_norm_correction(X, p.grad, channel_first=False)
             p.grad.set_(X.view_as(p)) # pyright:ignore[reportArgumentType]
@@ -100,7 +103,7 @@ class Orthogonalize(TensorTransform):
 
     standard Muon with Adam fallback
     ```py
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.head.parameters(),
         tz.m.Split(
             # apply muon only to 2D+ parameters
@@ -131,7 +134,7 @@ class Orthogonalize(TensorTransform):
 
         if not orthogonalize: return tensor
 
-        if _is_at_least_2d(tensor):
+        if _is_at_least_2d(tensor, channel_first=channel_first):
 
             X = _orthogonalize_format(tensor, method, channel_first=channel_first)
 
@@ -173,7 +176,7 @@ class MuonAdjustLR(Transform):
         alphas = [s['alpha'] for s in settings]
         channel_first = [s["channel_first=channel_first"] for s in settings]
         tensors_alphas = [
-            (t, adjust_lr_for_muon(a, t.shape, cf)) for t, a, cf in zip(tensors, alphas, channel_first) if _is_at_least_2d(t)
+            (t, adjust_lr_for_muon(a, t.shape, cf)) for t, a, cf in zip(tensors, alphas, channel_first) if _is_at_least_2d(t, channel_first=cf)
         ]
         tensors = [i[0] for i in tensors_alphas]
         a = [i[1] for i in alphas]

torchzero/modules/adaptive/natural_gradient.py

@@ -4,7 +4,7 @@ from ...core import Transform
 from ...utils.derivatives import jacobian_wrt, flatten_jacobian
 from ...utils import vec_to_tensors
 from ...linalg import linear_operator
-from .lmadagrad import lm_adagrad_apply, lm_adagrad_update
+from .ggt import ggt_update
 
 class NaturalGradient(Transform):
     """Natural gradient approximated via empirical fisher information matrix.
@@ -41,7 +41,7 @@ class NaturalGradient(Transform):
     y = torch.randn(64, 10)
 
     model = nn.Sequential(nn.Linear(20, 64), nn.ELU(), nn.Linear(64, 10))
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.NaturalGradient(),
         tz.m.LR(3e-2)
@@ -61,7 +61,7 @@ class NaturalGradient(Transform):
     y = torch.randn(64, 10)
 
     model = nn.Sequential(nn.Linear(20, 64), nn.ELU(), nn.Linear(64, 10))
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.NaturalGradient(),
         tz.m.LR(3e-2)
@@ -84,7 +84,7 @@ class NaturalGradient(Transform):
         return torch.stack([(1 - x1).abs(), (10 * (x2 - x1**2).abs())])
 
     X = torch.tensor([-1.1, 2.5], requires_grad=True)
-    opt = tz.Modular([X], tz.m.NaturalGradient(sqrt=True, gn_grad=True), tz.m.LR(0.05))
+    opt = tz.Optimizer([X], tz.m.NaturalGradient(sqrt=True, gn_grad=True), tz.m.LR(0.05))
 
     for iter in range(200):
         losses = rosenbrock(X)
@@ -99,18 +99,24 @@ class NaturalGradient(Transform):
     @torch.no_grad
     def update_states(self, objective, states, settings):
         params = objective.params
+        closure = objective.closure
         fs = settings[0]
         batched = fs['batched']
         gn_grad = fs['gn_grad']
 
-        closure = objective.closure
-        assert closure is not None
+        # compute per-sample losses
+        f = objective.loss
+        if f is None:
+            assert closure is not None
+            with torch.enable_grad():
+                f = objective.get_loss(backward=False) # n_out
+                assert isinstance(f, torch.Tensor)
 
+        # compute per-sample gradients
         with torch.enable_grad():
-            f = objective.get_loss(backward=False) # n_out
-            assert isinstance(f, torch.Tensor)
             G_list = jacobian_wrt([f.ravel()], params, batched=batched)
 
+        # set scalar loss and it's grad to objective
         objective.loss = f.sum()
         G = self.global_state["G"] = flatten_jacobian(G_list) # (n_samples, ndim)
 
@@ -123,8 +129,10 @@ class NaturalGradient(Transform):
         objective.grads = vec_to_tensors(g, params)
 
         # set closure to calculate scalar value for line searches etc
-        if objective.closure is not None:
+        if closure is not None:
+
             def ngd_closure(backward=True):
+
                 if backward:
                     objective.zero_grad()
                     with torch.enable_grad():
@@ -152,22 +160,31 @@ class NaturalGradient(Transform):
         if sqrt:
             # this computes U, S <- SVD(M), then calculate update as U S^-1 Uᵀg,
             # but it computes it through eigendecompotision
-            U, L = lm_adagrad_update(G.H, reg, 0)
-            if U is None or L is None: return objective
+            L, U = ggt_update(G.H, damping=reg, rdamping=1e-16, truncate=0, eig_tol=1e-12)
+
+            if U is None or L is None:
+
+                # fallback to element-wise
+                g = self.global_state["g"]
+                g /= G.square().mean(0).sqrt().add(reg)
+                objective.updates = vec_to_tensors(g, params)
+                return objective
 
-            v = lm_adagrad_apply(self.global_state["g"], U, L)
+            # whiten
+            z = U.T @ self.global_state["g"]
+            v = (U * L.rsqrt()) @ z
             objective.updates = vec_to_tensors(v, params)
             return objective
 
         # we need (G^T G)v = g
         # where g = G^T
         # so we need to solve (G^T G)v = G^T
-        GGT = G @ G.H # (n_samples, n_samples)
+        GGt = G @ G.H # (n_samples, n_samples)
 
         if reg != 0:
-            GGT.add_(torch.eye(GGT.size(0), device=GGT.device, dtype=GGT.dtype).mul_(reg))
+            GGt.add_(torch.eye(GGt.size(0), device=GGt.device, dtype=GGt.dtype).mul_(reg))
 
-        z, _ = torch.linalg.solve_ex(GGT, torch.ones_like(GGT[0])) # pylint:disable=not-callable
+        z, _ = torch.linalg.solve_ex(GGt, torch.ones_like(GGt[0])) # pylint:disable=not-callable
         v = G.H @ z
 
         objective.updates = vec_to_tensors(v, params)

torchzero/modules/adaptive/psgd/__init__.py

@@ -0,0 +1,5 @@
+from .psgd_dense_newton import PSGDDenseNewton
+from .psgd_kron_newton import PSGDKronNewton
+from .psgd_kron_whiten import PSGDKronWhiten
+from .psgd_lra_newton import PSGDLRANewton
+from .psgd_lra_whiten import PSGDLRAWhiten

torchzero/modules/adaptive/psgd/_psgd_utils.py

@@ -0,0 +1,37 @@
+# pylint:disable=not-callable
+import warnings
+
+import torch
+
+from .psgd import lift2single
+
+
+def _initialize_lra_state_(tensor: torch.Tensor, state, setting):
+    n = tensor.numel()
+    rank = max(min(setting["rank"], n-1), 1)
+    dtype=tensor.dtype
+    device=tensor.device
+
+    U = torch.randn((n, rank), dtype=dtype, device=device)
+    U *= 0.1**0.5 / torch.linalg.vector_norm(U)
+
+    V = torch.randn((n, rank), dtype=dtype, device=device)
+    V *= 0.1**0.5 / torch.linalg.vector_norm(V)
+
+    if setting["init_scale"] is None:
+        # warnings.warn("FYI: Will set the preconditioner initial scale on the fly. Recommend to set it manually.")
+        d = None
+    else:
+        d = torch.ones(n, 1, dtype=dtype, device=device) * setting["init_scale"]
+
+    state["UVd"] = [U, V, d]
+    state["Luvd"] = [lift2single(torch.zeros([], dtype=dtype, device=device)) for _ in range(3)]
+
+
+
+def _wrap_with_no_backward(opt):
+    """to use original psgd opts with visualbench"""
+    class _Wrapped:
+        def step(self, closure):
+            return opt.step(lambda: closure(False))
+    return _Wrapped()