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

torchzero/modules/adaptive/msam.py

@@ -2,9 +2,9 @@ from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, Target, Transform, apply_transform
+from ...core import Chainable, Module,  Transform, TensorTransform, step, Objective
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states, generic_ne
-from ..functional import ema_
+from ..opt_utils import ema_
 from ..momentum.momentum import nag_
 
 
@@ -21,7 +21,7 @@ def msam_(
 
     # inner args
     inner: Module | None = None,
-    grads: list[torch.Tensor] | None = None,
+    objective: Objective | None = None,
 ):
     # weights w and wh, momentum μ, perturbation strength ρ
     # w = wh + rho * v / ||v||
@@ -54,8 +54,8 @@ def msam_(
     v1n = velocity_ / denom
 
     if inner is not None:
-        assert params is not None
-        inner_update = TensorList(apply_transform(inner, tensors, params=params, grads=grads))
+        assert objective is not None and inner is not None
+        inner_update = TensorList(step(objective, inner).get_updates())
 
     else:
         assert lr is not None
@@ -69,7 +69,7 @@ def msam_(
 
     return update
 
-class MSAM(Transform):
+class MSAMMomentum(TensorTransform):
     """Momentum-SAM from https://arxiv.org/pdf/2401.12033.
 
     This implementation expresses the update rule as function of gradient. This way it can be used as a drop-in
@@ -93,46 +93,40 @@ class MSAM(Transform):
         lerp (bool, optional):
             whether to use linear interpolation, if True, this becomes similar to exponential moving average. Defaults to False.
 
-    Examples:
-        MSAM
+    ### Examples:
 
-        .. code-block:: python
+    MSAM
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.MSAM(1e-3)
-            )
+    ```python
 
-        Adam with MSAM instead of exponential average. Note that this is different from Adam_MSAM.
-        To make Adam_MSAM and such, use the :code:`tz.m.MSAMObjective` module.
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.MSAM(1e-3)
+    )
+    ```
 
-        .. code-block:: python
+    Adam with MSAM instead of exponential average. Note that this is different from Adam_MSAM.
+    To make Adam_MSAM and such, use the ``tz.m.MSAMObjective`` module.
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.RMSprop(0.999, inner=tz.m.MSAM(1e-3)),
-                tz.m.Debias(0.9, 0.999),
-            )
+    ```python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.RMSprop(0.999, inner=tz.m.MSAM(1e-3)),
+        tz.m.Debias(0.9, 0.999),
+    )
+    ```
     """
-    _USES_LR = True
+
     def __init__(self, lr: float, momentum:float=0.9, rho:float=0.3,  weight_decay:float=0, nesterov=False, lerp=False,):
-        defaults = dict(momentum=momentum,rho=rho, nesterov=nesterov, lerp=lerp, weight_decay=weight_decay)
-        if self._USES_LR: defaults['lr'] = lr
+        defaults = dict(lr = lr, momentum=momentum, rho=rho, nesterov=nesterov, lerp=lerp, weight_decay=weight_decay)
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
-        s = self.settings[params[0]]
-        lerp = s['lerp']
-        nesterov = s['nesterov']
+        fs = settings[0]
 
-        if self._USES_LR:
-            lr, momentum, rho, weight_decay = unpack_dicts(settings, 'lr','momentum','rho','weight_decay', cls=NumberList)
-
-        else:
-            lr=None
-            momentum,rho,weight_decay = unpack_dicts(settings, 'momentum','rho','weight_decay', cls=NumberList)
+        lr, momentum, rho, weight_decay = unpack_dicts(settings, 'lr','momentum','rho','weight_decay', cls=NumberList)
 
         return msam_(
             TensorList(tensors),
@@ -142,16 +136,16 @@ class MSAM(Transform):
             lr=lr,
             rho=rho,
             weight_decay=weight_decay,
-            nesterov=nesterov,
-            lerp=lerp,
+            nesterov=fs['nesterov'],
+            lerp=fs['lerp'],
 
             # inner args
-            inner=self.children.get("modules", None),
-            grads=grads,
+            inner=None,
+            objective=None,
         )
 
 
-class MSAMObjective(MSAM):
+class MSAM(Transform):
     """Momentum-SAM from https://arxiv.org/pdf/2401.12033.
 
     Note:
@@ -160,7 +154,7 @@ class MSAMObjective(MSAM):
         to an incorrect update rule.
 
     Args:
-        modules (Chainable): modules that will optimizer the MSAM objective. Make sure :code:`tz.m.LR` is one of them.
+        modules (Chainable): modules that will optimize the MSAM objective. Make sure ``tz.m.LR`` is one of them.
         momentum (float, optional): momentum (beta). Defaults to 0.9.
         rho (float, optional): perturbation strength. Defaults to 0.3.
         nesterov (bool, optional): whether to use nesterov momentum formula. Defaults to False.
@@ -169,20 +163,44 @@ class MSAMObjective(MSAM):
             Defaults to False.
 
     Examples:
-        AdamW-MSAM
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                bench.parameters(),
-                tz.m.MSAMObjective(
-                    [tz.m.Adam(), tz.m.WeightDecay(1e-3), tz.m.LR(1e-3)],
-                    rho=1.
-                )
-            )
+    AdamW-MSAM
+
+    ```py
+    opt = tz.Optimizer(
+        bench.parameters(),
+        tz.m.MSAMObjective(
+            [tz.m.Adam(), tz.m.WeightDecay(1e-3), tz.m.LR(1e-3)],
+            rho=1.
+        )
+    )
+    ```
     """
-    _USES_LR = False
     def __init__(self, modules: Chainable, momentum:float=0.9, rho:float=0.3, weight_decay:float=0, nesterov=False, lerp=False):
-        super().__init__(lr=0, momentum=momentum, rho=rho, weight_decay=weight_decay, nesterov=nesterov, lerp=lerp)
+        defaults = dict(momentum=momentum, rho=rho, weight_decay=weight_decay, nesterov=nesterov, lerp=lerp)
+        super().__init__(defaults)
+
         self.set_child('modules', modules)
 
+
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        velocity = unpack_states(states, objective.params, 'velocity', cls=TensorList)
+        fs = settings[0]
+
+        momentum, rho, weight_decay = unpack_dicts(settings, 'momentum', 'rho', 'weight_decay', cls=NumberList)
+
+        return msam_(
+            TensorList(objective.get_updates()),
+            params=TensorList(objective.params),
+            velocity_=velocity,
+            momentum=momentum,
+            lr=None,
+            rho=rho,
+            weight_decay=weight_decay,
+            nesterov=fs['nesterov'],
+            lerp=fs['lerp'],
+
+            # inner args
+            inner=self.children["modules"],
+            objective=objective,
+        )

torchzero/modules/adaptive/muon.py

@@ -1,152 +1,85 @@
 from operator import itemgetter
 import math
-import warnings
-from collections.abc import Iterable, Sequence
-from typing import Literal
+from collections.abc import Iterable
 
 import torch
 
-from ...core import Modular, TensorwiseTransform, Target, Transform
-from ...utils import enable_compilation
-
+from ...core import TensorTransform,  Transform
+from ...linalg.orthogonalize import orthogonalize as _orthogonalize, OrthogonalizeMethod
 
 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
 
-# stolen from:
-# https://github.com/KellerJordan/Muon/blob/master/muon.py
-# actually at this stage its a frankenstein
-@enable_compilation
-def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int) -> torch.Tensor:
-    """
-    Applies to last 2 dims - so usually reverse_dims should be applied to G before and after.
-
-    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
-    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
-    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
-    zero even beyond the point where the iteration no longer converges all the way to one everywhere
-    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
-    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
-    performance at all relative to UV^T, where USV^T = G is the SVD.
-    """
-    assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
-    a, b, c = (3.4445, -4.7750,  2.0315)
-    X = G.bfloat16()
-    if G.size(-2) > G.size(-1):
-        X = X.mT
-
-    # Ensure spectral norm is at most 1
-    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
-    # Perform the NS iterations
-    for _ in range(steps):
-        A = X @ X.mT
-        B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
-        X = a * X + B @ X
-
-    if G.size(-2) > G.size(-1):
-        X = X.mT
-    return X
-
-# stolen from https://github.com/MarkTuddenham/Orthogonal-Optimisers.
-# Tuddenham, M., Prügel-Bennett, A., & Hare, J. (2022).
-# Orthogonalising gradients to speed up neural network optimisation. arXiv preprint arXiv:2202.07052.
-@torch.no_grad
-def _svd_orthogonalize(G: torch.Tensor, warn_fail=True) -> torch.Tensor:
-    """
-    Applies to first 2 dims and isn't batched - rest of dimensions are flattened.
-    """
-    X = G.view(G.shape[0], -1)
-
-    t = False
-    if X.size(0) > X.size(1):
-        X = X.T
-        t = True
-
-    orth_X: torch.Tensor | None = None
-    try:
-        u, s, vt = torch.linalg.svd(X, full_matrices=False) # pylint:disable=not-callable
-        orth_X = u @ vt
-    except RuntimeError:
-        # if warn: logging.warning('Failed to perform SVD, adding some noise.')
-        try:
-            u, s, v = torch.svd_lowrank(
-                X,
-                q=1,    # assume rank is at least 1
-                M=1e-4 * X.mean() * torch.randn_like(X))
-            orth_X = u @ v.T
-        except RuntimeError:
-            if warn_fail: warnings.warn(('Failed to perform SVD with noise,'
-                            ' skipping gradient orthogonalisation'))
-    if orth_X is not None:
-        if t: orth_X = orth_X.T
-        return orth_X.view_as(G)
-
-    return G # fail
+def _orthogonalize_format(
+    tensor: torch.Tensor,
+    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))
 
+    return _orthogonalize(tensor, method=method)
 
 @torch.no_grad
-def _dual_norm_correction(X: torch.Tensor, g: torch.Tensor, batch_first):
-    """batch first means it applies to last 2 dims, otherwise to 1st two dims"""
+def _dual_norm_correction(X: torch.Tensor, g: torch.Tensor, channel_first: bool):
+    """``channel_first`` means it applies to first two dims, otherwise to last two dims"""
     # this is from https://github.com/leloykun/adaptive-muon
     # Adaptive scaling,`(G * X).sum() * X` == (G.T @ X).trace() * X
-    if batch_first: X = torch.einsum('...ij,...ij,...ab->...ab', g.type_as(X), X, X)
-    else: X = torch.einsum('ij...,ij...,ab...->ab...', g.type_as(X), X, X)
+    if channel_first: X = torch.einsum('ij...,ij...,ab...->ab...', g.type_as(X), X, X)
+    else: X = torch.einsum('...ij,...ij,...ab->...ab', g.type_as(X), X, X)
     return X
 
 
 # code from
 # https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py
-def adjust_lr_for_muon(lr, param_shape):
-    A, B = param_shape[:2]
+def adjust_lr_for_muon(lr, param_shape, channel_first:bool):
+    if channel_first: A, B = param_shape[:2]
+    else: A, B = param_shape[-2:]
+
     # We adjust the learning rate and weight decay based on the size of the parameter matrix
     # as describted in the paper
     adjusted_ratio = 0.2 * math.sqrt(max(A, B))
     adjusted_lr = lr * adjusted_ratio
     return adjusted_lr
 
-def _orthogonalize_tensor(
-    tensor: torch.Tensor,
-    steps: int = 5,
-    method: Literal["newton-schulz", "svd"] = "newton-schulz",
-):
-    if method == 'newton-schulz': return reverse_dims(zeropower_via_newtonschulz5(reverse_dims(tensor), steps)).type_as(tensor)
-    if method == 'svd': return _svd_orthogonalize(tensor, False)
-    raise ValueError(method)
-
 
 def orthogonalize_grads_(
     params: Iterable[torch.Tensor],
-    steps: int = 5,
     dual_norm_correction=False,
-    method: Literal["newton-schulz", "svd"] = "newton-schulz",
+    method: OrthogonalizeMethod = "newtonschulz",
+    channel_first:bool=True,
 ):
-    """Uses newton-Schulz iteration to compute the zeroth power / orthogonalization of gradients of an iterable of parameters.
+    """Computes the zeroth power / orthogonalization of gradients of an iterable of parameters.
 
     This sets gradients in-place. Applies along first 2 dims (expected to be `out_channels, in_channels`).
 
     Note that the Muon page says that embeddings and classifier heads should not be orthogonalized.
     Args:
         params (abc.Iterable[torch.Tensor]): parameters that hold gradients to orthogonalize.
-        steps (int, optional):
-            The number of Newton-Schulz iterations to run. Defaults to 5.
         dual_norm_correction (bool, optional):
             enables dual norm correction from https://github.com/leloykun/adaptive-muon. Defaults to False.
         method (str, optional):
             Newton-Schulz is very fast, SVD is extremely slow but can be slighly more precise.
+        channel_first (bool, optional):
+            if True, orthogonalizes along 1st two dimensions, otherwise along last 2. Other dimensions
+            are considered batch dimensions.
     """
     for p in params:
-        if (p.grad is not None) and _is_at_least_2d(p.grad):
-            X = _orthogonalize_tensor(p.grad, steps, method)
-            if dual_norm_correction: X = _dual_norm_correction(X, p.grad, batch_first=False)
+        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]
 
 
 
-class Orthogonalize(TensorwiseTransform):
+class Orthogonalize(TensorTransform):
     """Uses Newton-Schulz iteration or SVD to compute the zeroth power / orthogonalization of update along first 2 dims.
 
     To disable orthogonalization for a parameter, put it into a parameter group with "orthogonalize" = False.
@@ -156,22 +89,21 @@ class Orthogonalize(TensorwiseTransform):
     To make Muon, use Split with Adam on 1d params
 
     Args:
-        ns_steps (int, optional):
-            The number of Newton-Schulz iterations to run. Defaults to 5.
         adjust_lr (bool, optional):
             Enables LR adjustment based on parameter size from "Muon is Scalable for LLM Training". Defaults to False.
         dual_norm_correction (bool, optional):
             enables dual norm correction from https://github.com/leloykun/adaptive-muon. Defaults to False.
         method (str, optional):
-            Newton-Schulz is very fast, SVD is extremely slow but can be slighly more precise.
-        target (str, optional):
-            what to set on var.
+            Newton-Schulz is very fast, SVD is slow but can be more precise.
+        channel_first (bool, optional):
+            if True, orthogonalizes along 1st two dimensions, otherwise along last 2. Other dimensions
+            are considered batch dimensions.
 
     ## Examples:
 
     standard Muon with Adam fallback
     ```py
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.head.parameters(),
         tz.m.Split(
             # apply muon only to 2D+ parameters
@@ -190,56 +122,62 @@ class Orthogonalize(TensorwiseTransform):
     Reference:
         Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, Jeremy Bernstein - Muon: An optimizer for hidden layers in neural networks (2024) https://github.com/KellerJordan/Muon
     """
-    def __init__(self, ns_steps=5, adjust_lr=False, dual_norm_correction=False,
-                 method: Literal['newton-schulz', 'svd'] = 'newton-schulz', target:Target='update'):
-        defaults = dict(orthogonalize=True, ns_steps=ns_steps, dual_norm_correction=dual_norm_correction, adjust_lr=adjust_lr, method=method.lower())
-        super().__init__(uses_grad=False, defaults=defaults, target=target)
+    def __init__(self, adjust_lr=False, dual_norm_correction=False,
+                 method: OrthogonalizeMethod = 'newtonschulz', channel_first:bool=True):
+        defaults = dict(orthogonalize=True, dual_norm_correction=dual_norm_correction, adjust_lr=adjust_lr, method=method.lower(), channel_first=channel_first)
+        super().__init__(defaults=defaults)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, setting):
-        orthogonalize, ns_steps, dual_norm_correction, adjust_lr, method = itemgetter(
-            'orthogonalize', 'ns_steps', 'dual_norm_correction', 'adjust_lr', 'method')(setting)
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
+        orthogonalize, dual_norm_correction, adjust_lr, method, channel_first = itemgetter(
+            'orthogonalize', 'dual_norm_correction', 'adjust_lr', 'method', 'channel_first')(setting)
 
         if not orthogonalize: return tensor
 
-        if _is_at_least_2d(tensor):
+        if _is_at_least_2d(tensor, channel_first=channel_first):
 
-            X = _orthogonalize_tensor(tensor, ns_steps, method)
+            X = _orthogonalize_format(tensor, method, channel_first=channel_first)
 
             if dual_norm_correction:
-                X = _dual_norm_correction(X, tensor, batch_first=False)
+                X = _dual_norm_correction(X, tensor, channel_first=channel_first)
 
             if adjust_lr:
-                X.mul_(adjust_lr_for_muon(1, param.shape))
+                X.mul_(adjust_lr_for_muon(1, param.shape, channel_first=channel_first))
 
             return X.view_as(param)
 
         return tensor
 
 
-class DualNormCorrection(TensorwiseTransform):
+class DualNormCorrection(TensorTransform):
     """Dual norm correction for dualizer based optimizers (https://github.com/leloykun/adaptive-muon).
     Orthogonalize already has this built in with the `dual_norm_correction` setting."""
-    def __init__(self, target: Target='update'):
-        super().__init__({}, uses_grad=True, target=target)
+    def __init__(self, channel_first: bool = True):
+        defaults = dict(channel_first=channel_first)
+        super().__init__(defaults)
 
-    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+    @torch.no_grad
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
         assert grad is not None
         if (tensor.ndim >= 2) and (tensor.size(0) > 1) and (tensor.size(1) > 1):
-            return _dual_norm_correction(tensor, grad, batch_first=False)
+            return _dual_norm_correction(tensor, grad, channel_first=setting["channel_first"])
         return tensor
 
 
 class MuonAdjustLR(Transform):
     """LR adjustment for Muon from "Muon is Scalable for LLM Training" (https://github.com/MoonshotAI/Moonlight/tree/master).
-    Orthogonalize already has this built in with the `adjust_lr` setting, however you might want to move this to be later in the chain."""
-    def __init__(self, alpha: float = 1, target: Target='update'):
-        defaults = dict(alpha=alpha)
-        super().__init__(defaults=defaults, uses_grad=False, target=target)
+    Orthogonalize already has this built in with the ``adjust_lr`` setting, however you might want to move this to be later in the chain."""
+    def __init__(self, channel_first: bool = True, alpha: float = 1):
+        defaults = dict(channel_first=channel_first, alpha=alpha)
+        super().__init__(defaults=defaults)
 
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    @torch.no_grad
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         alphas = [s['alpha'] for s in settings]
-        tensors_alphas = [(t, adjust_lr_for_muon(a, t.shape)) for t, a in zip(tensors, alphas) if _is_at_least_2d(t)]
+        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, channel_first=cf)
+        ]
         tensors = [i[0] for i in tensors_alphas]
         a = [i[1] for i in alphas]
         torch._foreach_mul_(tensors, a)