Skip to content

Commit 8b0c3b0

Browse files
authored
feat(pt): add HybridMuonOptimizer (#5149)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added HybridMuon optimizer option with configurable hyperparameters (momentum alias, lr_adjust, lr_adjust_coeff, min_2d_dim, muon_2d_only, weight decay) and integrated training support (initialization, LR tracking, scheduler compatibility, checkpoint resume). * **Tests** * Added comprehensive tests covering orthogonalization, optimizer step behavior, weight-decay effects, fallback cases, LR-adjust modes, and state save/load. <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub> <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent 9553e6e commit 8b0c3b0

File tree

5 files changed

+1087
-13
lines changed

5 files changed

+1087
-13
lines changed

deepmd/pt/optimizer/__init__.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,19 @@
22
from .adamuon import (
33
AdaMuonOptimizer,
44
)
5+
from .hybrid_muon import (
6+
HybridMuonOptimizer,
7+
)
58
from .KFWrapper import (
69
KFOptimizerWrapper,
710
)
811
from .LKF import (
912
LKFOptimizer,
1013
)
1114

12-
__all__ = ["AdaMuonOptimizer", "KFOptimizerWrapper", "LKFOptimizer"]
15+
__all__ = [
16+
"AdaMuonOptimizer",
17+
"HybridMuonOptimizer",
18+
"KFOptimizerWrapper",
19+
"LKFOptimizer",
20+
]

0 commit comments

Comments
 (0)