fix: FSDP2 do not support foreach ops in HybridMuon

Author: OutisLi

deepmd/pt/optimizer/hybrid_muon.py

@@ -84,22 +84,7 @@
 NS_COEFF_C: float = 2.0315
 
 
-def _maybe_compile(
-    fn: callable,
-) -> callable:
-    """Compile a function if torch.compile is available."""
-    if not hasattr(torch, "compile"):
-        return fn
-    # Skip compile if default device is CUDA but CUDA is unavailable.
-    if hasattr(torch, "get_default_device"):
-        default_device = torch.get_default_device()
-        if default_device.type == "cuda" and not torch.cuda.is_available():
-            return fn
-    return torch.compile(fn, fullgraph=True, dynamic=True)
-
-
-@_maybe_compile
-def _zeropower_via_newtonschulz5_2d(
+def _newton_schulz_orth(
     G: torch.Tensor,
 ) -> torch.Tensor:
     """
@@ -132,70 +117,6 @@ def _zeropower_via_newtonschulz5_2d(
     return X
 
 
-@_maybe_compile
-def _zeropower_via_newtonschulz5_3d(
-    G: torch.Tensor,
-) -> torch.Tensor:
-    """
-    Orthogonalize a 3D batch of matrices via quintic Newton-Schulz iteration.
-
-    Mathematical formulation:
-        X_0 = G / ||G||_F
-        X_{k+1} = a*X_k + (b*A_k + c*A_k^2) @ X_k,  where A_k = X_k @ X_k^T
-        Coefficients: a=3.4445, b=-4.7750, c=2.0315
-    """
-    # === Step 1. Cast to bf16 and transpose tall matrices ===
-    X = G.to(dtype=torch.bfloat16)
-    transposed = X.size(-2) > X.size(-1)
-    if transposed:
-        X = X.transpose(-2, -1)
-
-    # === Step 2. Normalize Frobenius norm to at most 1 ===
-    X = X / X.norm(dim=(-2, -1), keepdim=True).clamp(min=EPS)
-
-    # === Step 3. Newton-Schulz iterations with batched fused GEMM ===
-    for _ in range(NS_STEPS):
-        A = torch.bmm(X, X.transpose(-2, -1))
-        gram_update = torch.baddbmm(A, A, A, beta=NS_COEFF_B, alpha=NS_COEFF_C)
-        X = torch.baddbmm(X, gram_update, X, beta=NS_COEFF_A, alpha=1.0)
-
-    # === Step 4. Transpose back if needed ===
-    if transposed:
-        X = X.transpose(-2, -1)
-
-    return X
-
-
-def zeropower_via_newtonschulz5(
-    G: torch.Tensor,
-) -> torch.Tensor:
-    """
-    Compute the zeroth power (orthogonalization) via Newton-Schulz iteration.
-
-    Dispatches to compiled 2D or 3D kernels for best performance.
-
-    Parameters
-    ----------
-    G : torch.Tensor
-        Input matrix with shape (M, N) or batched input with shape (B, M, N).
-
-    Returns
-    -------
-    torch.Tensor
-        Orthogonalized tensor in bfloat16 with same shape as input.
-
-    Raises
-    ------
-    ValueError
-        If input is not 2D or 3D.
-    """
-    if G.ndim == 2:
-        return _zeropower_via_newtonschulz5_2d(G)
-    if G.ndim == 3:
-        return _zeropower_via_newtonschulz5_3d(G)
-    raise ValueError("Input must be 2D or 3D for Newton-Schulz orthogonalization.")
-
-
 def should_fallback_to_adam_for_matrix(
     p: torch.Tensor,
     min_2d_dim: int,
@@ -478,9 +399,11 @@ def step(
 
                 # exp_avg = beta1 * exp_avg + (1 - beta1) * grad
                 # exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * grad^2
-                torch._foreach_lerp_(adam_exp_avgs, adam_grads_fp32, 1 - adam_betas[0])
-                grad_sq = torch._foreach_mul(adam_grads_fp32, adam_grads_fp32)
-                torch._foreach_lerp_(adam_exp_avg_sqs, grad_sq, 1 - adam_betas[1])
+                for ea, g in zip(adam_exp_avgs, adam_grads_fp32):
+                    ea.lerp_(g, 1 - adam_betas[0])
+                grad_sq = [g * g for g in adam_grads_fp32]
+                for eas, gsq in zip(adam_exp_avg_sqs, grad_sq):
+                    eas.lerp_(gsq, 1 - adam_betas[1])
 
                 # === Step 1.3. Bias correction and parameter update ===
                 for i, p in enumerate(adam_params):
@@ -531,11 +454,11 @@ def step(
 
                 # exp_avg = beta1 * exp_avg + (1 - beta1) * grad
                 # exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * grad^2
-                torch._foreach_lerp_(
-                    adam_nd_exp_avgs, adam_nd_grads_fp32, 1 - adam_betas[0]
-                )
-                grad_sq = torch._foreach_mul(adam_nd_grads_fp32, adam_nd_grads_fp32)
-                torch._foreach_lerp_(adam_nd_exp_avg_sqs, grad_sq, 1 - adam_betas[1])
+                for ea, g in zip(adam_nd_exp_avgs, adam_nd_grads_fp32):
+                    ea.lerp_(g, 1 - adam_betas[0])
+                grad_sq = [g * g for g in adam_nd_grads_fp32]
+                for eas, gsq in zip(adam_nd_exp_avg_sqs, grad_sq):
+                    eas.lerp_(gsq, 1 - adam_betas[1])
 
                 # === Step 2.3. Bias correction and parameter update ===
                 for i, p in enumerate(adam_nd_params):
@@ -589,15 +512,11 @@ def step(
 
                 # exp_avg = beta1 * exp_avg + (1 - beta1) * grad
                 # exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * grad^2
-                torch._foreach_lerp_(
-                    adam_matrix_exp_avgs, adam_matrix_grads_fp32, 1 - adam_betas[0]
-                )
-                grad_sq_m = torch._foreach_mul(
-                    adam_matrix_grads_fp32, adam_matrix_grads_fp32
-                )
-                torch._foreach_lerp_(
-                    adam_matrix_exp_avg_sqs, grad_sq_m, 1 - adam_betas[1]
-                )
+                for ea, g in zip(adam_matrix_exp_avgs, adam_matrix_grads_fp32):
+                    ea.lerp_(g, 1 - adam_betas[0])
+                grad_sq_m = [g * g for g in adam_matrix_grads_fp32]
+                for eas, gsq in zip(adam_matrix_exp_avg_sqs, grad_sq_m):
+                    eas.lerp_(gsq, 1 - adam_betas[1])
 
                 # === Step 3.3. Compute unclipped deltas ===
                 raw_deltas: list[torch.Tensor] = []
@@ -611,8 +530,8 @@ def step(
 
                 # === Step 3.4. Clip updates by relative norm and apply ===
                 max_rel_change = 0.05
-                p_norms = torch.stack(torch._foreach_norm(adam_matrix_params))
-                delta_norms = torch.stack(torch._foreach_norm(raw_deltas))
+                p_norms = torch.stack([p.norm() for p in adam_matrix_params])
+                delta_norms = torch.stack([d.norm() for d in raw_deltas])
                 floors = torch.tensor(
                     adam_matrix_abs_floor,
                     device=p_norms.device,
@@ -653,18 +572,21 @@ def step(
 
             # === Step 4.2. Apply weight decay (Muon path only) ===
             if weight_decay > 0 and muon_params_for_decay:
-                torch._foreach_mul_(muon_params_for_decay, 1.0 - lr * weight_decay)
+                for p in muon_params_for_decay:
+                    p.mul_(1.0 - lr * weight_decay)
 
             if not active_entries:
                 continue
 
             # === Step 4.3. Momentum update (Nesterov) ===
             # m_t = beta * m_{t-1} + (1 - beta) * g_t
-            torch._foreach_lerp_(muon_momentum_buffers, muon_grads, 1 - momentum)
+            for buf, g in zip(muon_momentum_buffers, muon_grads):
+                buf.lerp_(g, 1 - momentum)
             # update = beta * m_t + (1 - beta) * g_t
-            muon_updates = torch._foreach_lerp(
-                muon_grads, muon_momentum_buffers, momentum
-            )
+            muon_updates = [
+                torch.lerp(g, buf, momentum)
+                for g, buf in zip(muon_grads, muon_momentum_buffers)
+            ]
 
             # === Step 4.4. Bucket by shape/device/dtype for batched NS ===
             buckets: dict[
@@ -689,37 +611,16 @@ def step(
                 else:
                     scale = max(1.0, rows / cols) ** 0.5
 
-                if len(bucket_entries) == 1:
-                    entry, update_tensor = bucket_entries[0]
+                # Process each entry individually with _newton_schulz_orth.
+                # compatible with sharding propagation under FSDP2.
+                for entry, update_tensor in bucket_entries:
                     update_matrix = update_tensor.reshape(rows, cols)
                     if not update_matrix.is_contiguous():
                         update_matrix = update_matrix.contiguous()
 
-                    orth = _zeropower_via_newtonschulz5_2d(update_matrix)
+                    orth = _newton_schulz_orth(update_matrix)
                     orth.mul_(scale)
                     delta = orth.reshape(entry["param"].shape)
                     entry["param"].add_(delta, alpha=-lr)
-                    continue
-
-                matrices: list[torch.Tensor] = []
-                params: list[torch.Tensor] = []
-                orig_shapes: list[tuple[int, ...]] = []
-
-                for entry, update_tensor in bucket_entries:
-                    update_matrix = update_tensor.reshape(rows, cols)
-                    matrices.append(
-                        update_matrix
-                        if update_matrix.is_contiguous()
-                        else update_matrix.contiguous()
-                    )
-                    params.append(entry["param"])
-                    orig_shapes.append(entry["param"].shape)
-
-                stacked = torch.stack(matrices, dim=0)
-                orth = _zeropower_via_newtonschulz5_3d(stacked)
-                orth.mul_(scale)
-
-                for i, _ in enumerate(bucket_entries):
-                    params[i].add_(orth[i].reshape(orig_shapes[i]), alpha=-lr)
 
         return loss

source/tests/pt/test_hybrid_muon.py

@@ -5,7 +5,7 @@
 
 from deepmd.pt.optimizer.hybrid_muon import (
     HybridMuonOptimizer,
-    zeropower_via_newtonschulz5,
+    _newton_schulz_orth,
 )
 from deepmd.pt.utils import (
     env,
@@ -48,7 +48,7 @@ def test_orthogonalization(self) -> None:
         """Test that NS produces approximately orthogonal output."""
         torch.manual_seed(42)
         G = torch.randn(4, 4, dtype=torch.float32, device=self.device)
-        X = zeropower_via_newtonschulz5(G)
+        X = _newton_schulz_orth(G)
 
         # X @ X.T should be approximately identity
         # Note: NS uses bf16 internally, 5 iterations gives ~0.1-0.3 error
@@ -68,17 +68,17 @@ def test_orthogonalization(self) -> None:
     def test_shape_and_dtype(self) -> None:
         """Test that output preserves shape and returns bf16."""
         torch.manual_seed(42)
-        for shape in [(4, 4), (6, 4), (3, 4, 4)]:
+        for shape in [(4, 4), (6, 4)]:
             G = torch.randn(*shape, dtype=torch.float32, device=self.device)
-            X = zeropower_via_newtonschulz5(G)
+            X = _newton_schulz_orth(G)
             self.assertEqual(X.shape, G.shape)
             self.assertEqual(X.dtype, torch.bfloat16)
 
     def test_invalid_input(self) -> None:
-        """Test that <2D input raises ValueError."""
+        """Test that 1D input raises error."""
         G_1d = torch.randn(10, dtype=torch.float32, device=self.device)
-        with self.assertRaises(ValueError):
-            zeropower_via_newtonschulz5(G_1d)
+        with self.assertRaises((ValueError, RuntimeError, IndexError)):
+            _newton_schulz_orth(G_1d)
 
 
 @unittest.skipIf(not BF16_SUPPORTED, "bf16 matmul not supported on this device")