sort args

Author: OutisLi

deepmd/pt/optimizer/adamuon.py

@@ -160,7 +160,7 @@ class AdaMuonOptimizer(Optimizer):
     Key AdaMuon features:
     - Sign-stabilized orthogonal direction: Applies sign() before orthogonalization.
     - Per-element second-moment normalization using momentum coefficient.
-    - RMS-aligned global scaling: 0.2 * sqrt(min * max) / norm.
+    - RMS-aligned global scaling: 0.2 * sqrt(m * n) / norm.
 
     Parameters
     ----------
@@ -245,7 +245,7 @@ def step(
 
         Returns
         -------
-        loss : float, optional
+        loss : torch.Tensor, optional
             The loss value if closure is provided.
         """
         loss = None

deepmd/pt/train/training.py

@@ -166,11 +166,9 @@ def get_opt_param(params: dict[str, Any]) -> tuple[str, dict[str, Any]]:
                 # Common parameters
                 "weight_decay": params.get("weight_decay", 0.001),
                 # Muon/AdaMuon parameters
-                "muon_momentum": params.get("muon_momentum", 0.95),
+                "momentum": params.get("momentum", 0.95),
                 "adam_beta1": params.get("adam_beta1", 0.9),
                 "adam_beta2": params.get("adam_beta2", 0.95),
-                "adam_eps": params.get("adam_eps", 1e-7),
-                "nesterov": params.get("nesterov", True),
             }
             return opt_type, opt_param
 
@@ -710,7 +708,7 @@ def warm_up_linear(step: int, warmup_steps: int) -> float:
             self.optimizer = AdaMuonOptimizer(
                 self.wrapper.parameters(),
                 lr=self.lr_exp.start_lr,
-                momentum=float(self.opt_param.get("muon_momentum", 0.95)),
+                momentum=float(self.opt_param.get("momentum", 0.95)),
                 weight_decay=float(self.opt_param.get("weight_decay", 0.001)),
                 adam_betas=(
                     float(self.opt_param.get("adam_beta1", 0.9)),

deepmd/utils/argcheck.py

@@ -3375,7 +3375,7 @@ def training_args(
                     dict,
                     [
                         Argument(
-                            "muon_momentum",
+                            "momentum",
                             float,
                             optional=True,
                             default=0.95,
@@ -3413,8 +3413,8 @@ def training_args(
                             default=10.0,
                             doc=doc_only_pt_supported
                             + "Learning rate adjustment factor for Adam (1D params). "
-                            "If lr_adjust <= 0: use match-RMS scaling for AdaMuon, Adam uses lr directly. "
-                            "If lr_adjust > 0: use rectangular correction for AdaMuon, Adam uses lr/lr_adjust.",
+                            "If lr_adjust <= 0: use match-RMS scaling (scale = lr_adjust_coeff * sqrt(max(m, n))), Adam uses lr directly. "
+                            "If lr_adjust > 0: use rectangular correction (scale = sqrt(max(1.0, m/n))), Adam uses lr/lr_adjust.",
                         ),
                         Argument(
                             "lr_adjust_coeff",