rename custom muon to hybridmuon

OutisLi · OutisLi · commit 1978c7f39a30 · 2026-01-13T10:48:04.000+08:00
diff --git a/deepmd/pt/optimizer/__init__.py b/deepmd/pt/optimizer/__init__.py
@@ -2,14 +2,19 @@
 from .adamuon import (
     AdaMuonOptimizer,
 )
+from .hybrid_muon import (
+    HybridMuonOptimizer,
+)
 from .KFWrapper import (
     KFOptimizerWrapper,
 )
 from .LKF import (
     LKFOptimizer,
 )
-from .muon import (
-    MuonOptimizer,
-)
 
-__all__ = ["AdaMuonOptimizer", "KFOptimizerWrapper", "LKFOptimizer", "MuonOptimizer"]
+__all__ = [
+    "AdaMuonOptimizer",
+    "HybridMuonOptimizer",
+    "KFOptimizerWrapper",
+    "LKFOptimizer",
+]
diff --git a/deepmd/pt/optimizer/hybrid_muon.py b/deepmd/pt/optimizer/hybrid_muon.py
@@ -1,10 +1,15 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """
-Muon optimizer for DeePMD-kit PyTorch backend.
+HybridMuon optimizer for DeePMD-kit PyTorch backend.
 
-Muon is an optimizer that applies Newton-Schulz orthogonalization to the gradient
-before using momentum, resulting in orthogonalized updates for weight matrices.
-This can improve training stability and convergence for certain architectures.
+HybridMuon is a HYBRID optimizer that automatically combines Muon and Adam:
+- For >=2D parameters with min(m,n) >= min_2d_dim: Muon update with Newton-Schulz
+- For 2D parameters with min(m,n) < min_2d_dim: Adam fallback with update clipping
+- For 1D parameters (biases, layer norms): Standard Adam
+
+This is different from PyTorch's torch.optim.Muon, which ONLY supports 2D parameters
+and requires manual configuration of AdamW for 1D parameters. HybridMuon provides
+automatic routing based on parameter dimensionality.
 
 Algorithm
 ---------
@@ -33,9 +38,15 @@
 - Muon gradients: cast to parameter dtype before momentum update
 - Adam gradients: cast to float32 for update computation
 
-Reference
----------
-https://github.com/KellerJordan/Muon
+References
+----------
+.. [1] Keller Jordan, "Muon: An optimizer for hidden layers in neural networks."
+       https://kellerjordan.github.io/posts/muon/
+       https://github.com/KellerJordan/Muon
+.. [2] Moonshot team, "Muon is Scalable for LLM Training," arXiv:2502.16982, 2025.
+       https://arxiv.org/abs/2502.16982
+.. [3] Moonlight GitHub Repository.
+       https://github.com/MoonshotAI/Moonlight
 """
 
 from __future__ import (
@@ -223,9 +234,9 @@ def should_fallback_to_adam_for_matrix(
     return min(m, n) < min_2d_dim
 
 
-class MuonOptimizer(Optimizer):
+class HybridMuonOptimizer(Optimizer):
     """
-    Muon optimizer with small-2D Adam fallback and 1D Adam path.
+    HybridMuon optimizer with small-2D Adam fallback and 1D Adam path.
 
     This optimizer applies different update rules based on parameter dimensionality:
     - For >=2D parameters with min(m, n) >= min_2d_dim:
@@ -286,7 +297,7 @@ class MuonOptimizer(Optimizer):
 
     Examples
     --------
-    >>> optimizer = MuonOptimizer(model.parameters(), lr=1e-3)
+    >>> optimizer = HybridMuonOptimizer(model.parameters(), lr=1e-3)
     >>> for epoch in range(epochs):
     ...     optimizer.zero_grad()
     ...     loss.backward()
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -43,9 +43,9 @@
 )
 from deepmd.pt.optimizer import (
     AdaMuonOptimizer,
+    HybridMuonOptimizer,
     KFOptimizerWrapper,
     LKFOptimizer,
-    MuonOptimizer,
 )
 from deepmd.pt.train.wrapper import (
     ModelWrapper,
@@ -730,8 +730,8 @@ def warm_up_linear(step: int, warmup_steps: int) -> float:
                 lr_adjust=float(self.opt_param.get("lr_adjust", 10.0)),
                 lr_adjust_coeff=float(self.opt_param.get("lr_adjust_coeff", 0.2)),
             )
-        elif self.opt_type == "Muon":
-            self.optimizer = MuonOptimizer(
+        elif self.opt_type == "HybridMuon":
+            self.optimizer = HybridMuonOptimizer(
                 self.wrapper.parameters(),
                 lr=self.lr_exp.start_lr,
                 momentum=float(self.opt_param.get("momentum", 0.95)),
@@ -820,7 +820,7 @@ def step(_step_id: int, task_key: str = "Default") -> None:
                 print_str = f"Step {_step_id}: sample system{log_dict['sid']}  frame{log_dict['fid']}\n"
                 fout1.write(print_str)
                 fout1.flush()
-            if self.opt_type in ["Adam", "AdamW", "AdaMuon", "Muon"]:
+            if self.opt_type in ["Adam", "AdamW", "AdaMuon", "HybridMuon"]:
                 cur_lr = self.scheduler.get_last_lr()[0]
                 if _step_id < self.warmup_steps:
                     pref_lr = _lr.start_lr
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
@@ -3452,7 +3452,7 @@ def training_args(
                     optional=True,
                 ),
                 Argument(
-                    "Muon",
+                    "HybridMuon",
                     dict,
                     [
                         Argument(
@@ -3462,7 +3462,7 @@ def training_args(
                             default=0.95,
                             alias=["muon_momentum"],
                             doc=doc_only_pt_supported
-                            + "Momentum coefficient for Muon optimizer (>=2D params). "
+                            + "Momentum coefficient for HybridMuon optimizer (>=2D params). "
                             "Used in Nesterov momentum update: m_t = beta*m_{t-1} + (1-beta)*g_t.",
                         ),
                         Argument(
@@ -3487,15 +3487,15 @@ def training_args(
                             optional=True,
                             default=0.001,
                             doc=doc_only_pt_supported
-                            + "Weight decay coefficient. Applied only to >=2D parameters (Muon path).",
+                            + "Weight decay coefficient. Applied only to >=2D parameters (HybridMuon path).",
                         ),
                         Argument(
                             "lr_adjust",
                             float,
                             optional=True,
                             default=10.0,
                             doc=doc_only_pt_supported
-                            + "Learning rate adjustment mode for Muon scaling and Adam learning rate. "
+                            + "Learning rate adjustment mode for HybridMuon scaling and Adam learning rate. "
                             "If lr_adjust <= 0: use match-RMS scaling (scale = coeff*sqrt(max(m,n))), Adam uses lr directly. "
                             "If lr_adjust > 0: use rectangular correction (scale = sqrt(max(1, m/n))), Adam uses lr/lr_adjust. "
                             "Default is 10.0 (Adam lr = lr/10).",
@@ -3515,14 +3515,20 @@ def training_args(
                             default=1,
                             alias=["muon_min_2d_dim"],
                             doc=doc_only_pt_supported
-                            + "Minimum min(m, n) threshold for Muon on 2D matrices. "
-                            "Matrices with min(m, n) >= min_2d_dim use Muon; "
+                            + "Minimum min(m, n) threshold for HybridMuon on 2D matrices. "
+                            "Matrices with min(m, n) >= min_2d_dim use HybridMuon; "
                             "those with min(m, n) < min_2d_dim use Adam fallback. "
                             "Set to 1 to disable fallback.",
                         ),
                     ],
                     [],
                     optional=True,
+                    doc=doc_only_pt_supported
+                    + "HybridMuon optimizer (DeePMD-kit custom implementation). "
+                    + "This is a Hybrid optimizer that automatically combines Muon and Adam. "
+                    + "For >=2D params: Muon update with Newton-Schulz. "
+                    + "For 1D params: Standard Adam. "
+                    + "This is DIFFERENT from PyTorch's torch.optim.Muon which ONLY supports 2D parameters.",
                 ),
             ],
             optional=True,
diff --git a/source/tests/pt/test_hybrid_muon.py b/source/tests/pt/test_hybrid_muon.py
@@ -3,8 +3,8 @@
 
 import torch
 
-from deepmd.pt.optimizer.muon import (
-    MuonOptimizer,
+from deepmd.pt.optimizer.hybrid_muon import (
+    HybridMuonOptimizer,
     zeropower_via_newtonschulz5,
 )
 from deepmd.pt.utils import (
@@ -82,8 +82,8 @@ def test_invalid_input(self) -> None:
 
 
 @unittest.skipIf(not BF16_SUPPORTED, "bf16 matmul not supported on this device")
-class TestMuonOptimizer(unittest.TestCase):
-    """Test MuonOptimizer class."""
+class TestHybridMuonOptimizer(unittest.TestCase):
+    """Test HybridMuonOptimizer class."""
 
     def setUp(self) -> None:
         self.device = env.DEVICE
@@ -96,7 +96,7 @@ def test_step(self) -> None:
             torch.nn.ReLU(),
             torch.nn.Linear(20, 5, device=self.device),
         )
-        optimizer = MuonOptimizer(model.parameters(), lr=0.02)
+        optimizer = HybridMuonOptimizer(model.parameters(), lr=0.02)
 
         x = torch.randn(4, 10, device=self.device)
         model(x).sum().backward()
@@ -111,7 +111,7 @@ def test_weight_decay(self) -> None:
         """Test weight decay reduces parameter norm."""
         torch.manual_seed(42)
         model = torch.nn.Linear(10, 10, device=self.device)
-        optimizer = MuonOptimizer(model.parameters(), lr=0.02, weight_decay=0.1)
+        optimizer = HybridMuonOptimizer(model.parameters(), lr=0.02, weight_decay=0.1)
 
         initial_norm = model.weight.norm().item()
         for _ in range(10):
@@ -126,7 +126,7 @@ def test_muon_adam_separation(self) -> None:
         """Test Muon for 2D params, Adam for 1D params."""
         torch.manual_seed(42)
         model = torch.nn.Linear(10, 10, device=self.device)
-        optimizer = MuonOptimizer(model.parameters(), lr=0.02)
+        optimizer = HybridMuonOptimizer(model.parameters(), lr=0.02)
 
         x = torch.randn(4, 10, device=self.device)
         model(x).sum().backward()
@@ -145,7 +145,7 @@ def test_muon_adam_fallback_small_2d(self) -> None:
         torch.manual_seed(42)
         linear_small = torch.nn.Linear(10, 1, bias=False, device=self.device)
         linear_large = torch.nn.Linear(10, 10, bias=False, device=self.device)
-        optimizer = MuonOptimizer(
+        optimizer = HybridMuonOptimizer(
             list(linear_small.parameters()) + list(linear_large.parameters()),
             lr=0.02,
             min_2d_dim=2,
@@ -172,8 +172,8 @@ def test_lr_adjust_modes(self) -> None:
         model2 = torch.nn.Linear(10, 20, bias=False, device=self.device)
         model2.load_state_dict(model1.state_dict())
 
-        opt1 = MuonOptimizer(model1.parameters(), lr=0.02, lr_adjust=0.0)
-        opt2 = MuonOptimizer(model2.parameters(), lr=0.02, lr_adjust=10.0)
+        opt1 = HybridMuonOptimizer(model1.parameters(), lr=0.02, lr_adjust=0.0)
+        opt2 = HybridMuonOptimizer(model2.parameters(), lr=0.02, lr_adjust=10.0)
 
         x = torch.randn(4, 10, device=self.device)
 
@@ -192,7 +192,7 @@ def test_lr_adjust_modes(self) -> None:
 
 
 @unittest.skipIf(not BF16_SUPPORTED, "bf16 matmul not supported on this device")
-class TestMuonOptimizerStateDict(unittest.TestCase):
+class TestHybridMuonOptimizerStateDict(unittest.TestCase):
     """Test optimizer state dict save/load."""
 
     def setUp(self) -> None:
@@ -202,7 +202,7 @@ def test_state_dict_save_load(self) -> None:
         """Test saving and loading optimizer state."""
         torch.manual_seed(42)
         model = torch.nn.Linear(10, 10, device=self.device)
-        optimizer = MuonOptimizer(model.parameters(), lr=0.02)
+        optimizer = HybridMuonOptimizer(model.parameters(), lr=0.02)
 
         for _ in range(3):
             optimizer.zero_grad()
@@ -212,7 +212,7 @@ def test_state_dict_save_load(self) -> None:
 
         state_dict = optimizer.state_dict()
 
-        optimizer2 = MuonOptimizer(model.parameters(), lr=0.02)
+        optimizer2 = HybridMuonOptimizer(model.parameters(), lr=0.02)
         optimizer2.load_state_dict(state_dict)
 
         # Verify state matches by param id, not iteration order
