feat(pt): Update AdaMuon

(cherry picked from commit ea5ac54)

Author: OutisLi

deepmd/pt/optimizer/adamuon.py

@@ -40,7 +40,7 @@
 def zeropower_via_newtonschulz5(
     G: torch.Tensor,
     steps: int = 5,
-    eps: float = 1e-7,
+    eps: float = 1e-8,
 ) -> torch.Tensor:
     """
     Compute the zeroth power (orthogonalization) of a matrix via Newton-Schulz iteration.
@@ -58,7 +58,7 @@ def zeropower_via_newtonschulz5(
     steps : int
         Number of Newton-Schulz iterations with default 5.
     eps : float
-        Numerical stability epsilon for norm clamping with default 1e-7.
+        Numerical stability epsilon for norm clamping with default 1e-8.
 
     Returns
     -------
@@ -177,7 +177,7 @@ class AdaMuonOptimizer(Optimizer):
     adam_betas : tuple[float, float]
         Adam beta coefficients with default (0.9, 0.95).
     adam_eps : float
-        Adam epsilon with default 1e-7.
+        Adam epsilon with default 1e-8.
     nesterov : bool
         Whether to use Nesterov momentum for AdaMuon with default True.
     lr_adjust : float
@@ -210,7 +210,7 @@ def __init__(
         weight_decay: float = 0.001,
         ns_steps: int = 5,
         adam_betas: tuple[float, float] = (0.9, 0.95),
-        adam_eps: float = 1e-7,
+        adam_eps: float = 1e-8,
         nesterov: bool = True,
         lr_adjust: float = 10.0,
         lr_adjust_coeff: float = 0.2,

deepmd/pt/train/training.py

@@ -711,7 +711,7 @@ def warm_up_linear(step: int, warmup_steps: int) -> float:
                 self.wrapper.parameters(),
                 lr=self.lr_exp.start_lr,
                 momentum=float(self.opt_param.get("muon_momentum", 0.95)),
-                weight_decay=float(self.opt_param.get("weight_decay", 0.0)),
+                weight_decay=float(self.opt_param.get("weight_decay", 0.001)),
                 adam_betas=(
                     float(self.opt_param.get("adam_beta1", 0.9)),
                     float(self.opt_param.get("adam_beta2", 0.95)),