Add DTensor compatible NS impl for Muon

Author: rwightman

timm/optim/muon.py

@@ -25,6 +25,11 @@
 from typing import List, Mapping, Optional, Sequence, Tuple, Union
 
 import torch
+try:
+    from torch.distributed.tensor import DTensor
+    has_dtensor = True
+except ImportError:
+    has_dtensor = False
 
 from ._types import ParamsT
 from .adamw import adamw
@@ -145,7 +150,9 @@ def zeropower_via_newtonschulz(
     if scale_eps:
         eps = scale_eps_for_ns(eps, G.shape)
 
-    X = G.to(dtype=dtype, copy=True)
+    is_dtensor = has_dtensor and isinstance(G, DTensor)
+
+    X = G.to(dtype=dtype) if is_dtensor else G.to(dtype=dtype, copy=True)
 
     # Transpose if needed (operate on dimension with fewer elements)
     transposed = X.size(-2) > X.size(-1)
@@ -159,21 +166,30 @@ def zeropower_via_newtonschulz(
     else:
         X.div_(X.norm(2, dim=(-2, -1), keepdim=True).mul(safety_factor).clamp_(min=eps))
 
-    # Batched vs unbatched fused MM
-    mm_fn = torch.baddbmm if X.ndim > 2 else torch.addmm
-
-    # Pre-allocate
-    X = X.contiguous()
-    A = torch.empty((*X.shape[:-1], X.size(-2)), device=X.device, dtype=X.dtype)
-    B = torch.empty_like(A)
-    C = torch.empty_like(X)
-
-    # Perform Newton-Schulz iterations
-    for a, b, c in coeff_sequence:
-        mm_fn(A, X, X.mT, beta=0.0, alpha=1.0, out=A)  # A = X @ X.mT
-        mm_fn(A, A, A, beta=b, alpha=c, out=B)  # B = b * A + c * A @ A
-        mm_fn(X, B, X, beta=a, alpha=1.0, out=C)  # C = a * X + B @ X
-        X, C = C, X  # swap refs to avoid copy
+    if is_dtensor:
+        # Basic, DTensor-friendly Newton-Schulz
+        for a, b, c in coeff_sequence:
+            A = X @ X.mT
+            B = b * A + c * (A @ A)
+            X = a * X + (B @ X)
+    else:
+        # Fast prealloc/out= path
+
+        # Batched vs unbatched fused MM
+        mm_fn = torch.baddbmm if X.ndim > 2 else torch.addmm
+
+        # Pre-allocate
+        X = X.contiguous()
+        A = torch.empty((*X.shape[:-1], X.size(-2)), device=X.device, dtype=X.dtype)
+        B = torch.empty_like(A)
+        C = torch.empty_like(X)
+
+        # Perform Newton-Schulz iterations
+        for a, b, c in coeff_sequence:
+            mm_fn(A, X, X.mT, beta=0.0, alpha=1.0, out=A)  # A = X @ X.mT
+            mm_fn(A, A, A, beta=b, alpha=c, out=B)  # B = b * A + c * A @ A
+            mm_fn(X, B, X, beta=a, alpha=1.0, out=C)  # C = a * X + B @ X
+            X, C = C, X  # swap refs to avoid copy
 
     if transposed:
         X = X.mT