Repository: kozistr/pytorch_optimizer (python)

Issue:
Add AdaMuon
## Paper or Code

https://arxiv.org/abs/2507.11005v1

Pull request:
[Feature] Implement `AdaMuon` optimizer
## Problem (Why?)

Close #394 

## Solution (What/How?)

- [x] implement `AdaMuon` optimizer

## Other changes (bug fixes, small refactors)

N/A

## Notes

N/A

Patch:
diff --git a/README.md b/README.md
index db32bbd0..3cc91d66 100644
--- a/README.md
+++ b/README.md
@@ -10,7 +10,7 @@
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **108 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Wide range of supported optimizers. Currently, **109 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
 * Including many variants such as `ADOPT`, `Cautious`, `AdamD`, `StableAdamW`, and `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
@@ -215,7 +215,8 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | RACS & Alice        | *Towards Efficient Optimizer Design for LLM via Structured Fisher Approximation with a Low-Rank Extension*     |                                                                                                                | <https://arxiv.org/pdf/2502.07752>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250207752G/exportcitation)                                                        |
 | VSGD                | *Variational Stochastic Gradient Descent for Deep Neural Networks*                                             | [github](https://github.com/generativeai-tue/vsgd)                                                             | <https://openreview.net/forum?id=xu4ATNjcdy>                                                | [cite](https://github.com/generativeai-tue/vsgd/tree/main?tab=readme-ov-file#cite)                                                  |
 | SNSM                | *Subset-Norm and Subspace-Momentum: Faster Memory-Efficient Adaptive Optimization with Convergence Guarantees* | [github](https://github.com/timmytonga/sn-sm)                                                                  | <https://arxiv.org/abs/2411.07120>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv241107120N/exportcitation)                                                        |
-| AdamC               | Why Gradients Rapidly Increase Near the End of Training*                                                       |                                                                                                                | <https://arxiv.org/abs/2506.02285>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250602285D/exportcitation)                                                        |
+| AdamC               | *Why Gradients Rapidly Increase Near the End of Training*                                                      |                                                                                                                | <https://arxiv.org/abs/2506.02285>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250602285D/exportcitation)                                                        |
+| AdaMuon             | *Adaptive Muon Optimizer*                                                                                      |                                                                                                                | <https://arxiv.org/abs/2507.11005v1>                                                        | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250711005S/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 
diff --git a/docs/changelogs/v3.6.2.md b/docs/changelogs/v3.6.2.md
new file mode 100644
index 00000000..75182055
--- /dev/null
+++ b/docs/changelogs/v3.6.2.md
@@ -0,0 +1,6 @@
+## Change Log
+
+### Feature
+
+* Implement `AdaMuon` optimizer. (#394, #395)
+    * [Adaptive Muon Optimizer](https://arxiv.org/abs/2507.11005v1)
diff --git a/docs/optimizer.md b/docs/optimizer.md
index fd969f8c..ac6a640f 100644
--- a/docs/optimizer.md
+++ b/docs/optimizer.md
@@ -284,6 +284,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.AdaMuon
+    :docstring:
+    :members:
+
 ::: pytorch_optimizer.Nero
     :docstring:
     :members:
diff --git a/poetry.lock b/poetry.lock
index 850f899b..77f30299 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -395,15 +395,15 @@ tqdm = ["tqdm"]
 
 [[package]]
 name = "fsspec"
-version = "2025.5.1"
+version = "2025.7.0"
 description = "File-system specification"
 optional = false
 python-versions = ">=3.9"
 groups = ["main"]
 markers = "python_version >= \"3.9\""
 files = [
-    {file = "fsspec-2025.5.1-py3-none-any.whl", hash = "sha256:24d3a2e663d5fc735ab256263c4075f374a174c3410c0b25e5bd1970bceaa462"},
-    {file = "fsspec-2025.5.1.tar.gz", hash = "sha256:2e55e47a540b91843b755e83ded97c6e897fa0942b11490113f09e9c443c2475"},
+    {file = "fsspec-2025.7.0-py3-none-any.whl", hash = "sha256:8b012e39f63c7d5f10474de957f3ab793b47b45ae7d39f2fb735f8bbe25c0e21"},
+    {file = "fsspec-2025.7.0.tar.gz", hash = "sha256:786120687ffa54b8283d942929540d8bc5ccfa820deb555a2b5d0ed2b737bf58"},
 ]
 
 [package.extras]
@@ -411,7 +411,7 @@ abfs = ["adlfs"]
 adl = ["adlfs"]
 arrow = ["pyarrow (>=1)"]
 dask = ["dask", "distributed"]
-dev = ["pre-commit", "ruff"]
+dev = ["pre-commit", "ruff (>=0.5)"]
 doc = ["numpydoc", "sphinx", "sphinx-design", "sphinx-rtd-theme", "yarl"]
 dropbox = ["dropbox", "dropboxdrivefs", "requests"]
 full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "dask", "distributed", "dropbox", "dropboxdrivefs", "fusepy", "gcsfs", "libarchive-c", "ocifs", "panel", "paramiko", "pyarrow (>=1)", "pygit2", "requests", "s3fs", "smbprotocol", "tqdm"]
@@ -431,7 +431,7 @@ smb = ["smbprotocol"]
 ssh = ["paramiko"]
 test = ["aiohttp (!=4.0.0a0,!=4.0.0a1)", "numpy", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "requests"]
 test-downstream = ["aiobotocore (>=2.5.4,<3.0.0)", "dask[dataframe,test]", "moto[server] (>4,<5)", "pytest-timeout", "xarray"]
-test-full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "cloudpickle", "dask", "distributed", "dropbox", "dropboxdrivefs", "fastparquet", "fusepy", "gcsfs", "jinja2", "kerchunk", "libarchive-c", "lz4", "notebook", "numpy", "ocifs", "pandas", "panel", "paramiko", "pyarrow", "pyarrow (>=1)", "pyftpdlib", "pygit2", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "python-snappy", "requests", "smbprotocol", "tqdm", "urllib3", "zarr", "zstandard"]
+test-full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "cloudpickle", "dask", "distributed", "dropbox", "dropboxdrivefs", "fastparquet", "fusepy", "gcsfs", "jinja2", "kerchunk", "libarchive-c", "lz4", "notebook", "numpy", "ocifs", "pandas", "panel", "paramiko", "pyarrow", "pyarrow (>=1)", "pyftpdlib", "pygit2", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "python-snappy", "requests", "smbprotocol", "tqdm", "urllib3", "zarr", "zstandard ; python_version < \"3.14\""]
 tqdm = ["tqdm"]
 
 [[package]]
@@ -1007,30 +1007,30 @@ testing = ["fields", "hunter", "process-tests", "pytest-xdist", "virtualenv"]
 
 [[package]]
 name = "ruff"
-version = "0.12.2"
+version = "0.12.4"
 description = "An extremely fast Python linter and code formatter, written in Rust."
 optional = false
 python-versions = ">=3.7"
 groups = ["dev"]
 files = [
-    {file = "ruff-0.12.2-py3-none-linux_armv6l.whl", hash = "sha256:093ea2b221df1d2b8e7ad92fc6ffdca40a2cb10d8564477a987b44fd4008a7be"},
-    {file = "ruff-0.12.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:09e4cf27cc10f96b1708100fa851e0daf21767e9709e1649175355280e0d950e"},
-    {file = "ruff-0.12.2-py3-none-macosx_11_0_arm64.whl", hash = "sha256:8ae64755b22f4ff85e9c52d1f82644abd0b6b6b6deedceb74bd71f35c24044cc"},
-    {file = "ruff-0.12.2-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3eb3a6b2db4d6e2c77e682f0b988d4d61aff06860158fdb413118ca133d57922"},
-    {file = "ruff-0.12.2-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:73448de992d05517170fc37169cbca857dfeaeaa8c2b9be494d7bcb0d36c8f4b"},
-    {file = "ruff-0.12.2-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3b8b94317cbc2ae4a2771af641739f933934b03555e51515e6e021c64441532d"},
-    {file = "ruff-0.12.2-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:45fc42c3bf1d30d2008023a0a9a0cfb06bf9835b147f11fe0679f21ae86d34b1"},
-    {file = "ruff-0.12.2-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ce48f675c394c37e958bf229fb5c1e843e20945a6d962cf3ea20b7a107dcd9f4"},
-    {file = "ruff-0.12.2-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:793d8859445ea47591272021a81391350205a4af65a9392401f418a95dfb75c9"},
-    {file = "ruff-0.12.2-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6932323db80484dda89153da3d8e58164d01d6da86857c79f1961934354992da"},
-    {file = "ruff-0.12.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:6aa7e623a3a11538108f61e859ebf016c4f14a7e6e4eba1980190cacb57714ce"},
-    {file = "ruff-0.12.2-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:2a4a20aeed74671b2def096bdf2eac610c7d8ffcbf4fb0e627c06947a1d7078d"},
-    {file = "ruff-0.12.2-py3-none-musllinux_1_2_i686.whl", hash = "sha256:71a4c550195612f486c9d1f2b045a600aeba851b298c667807ae933478fcef04"},
-    {file = "ruff-0.12.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:4987b8f4ceadf597c927beee65a5eaf994c6e2b631df963f86d8ad1bdea99342"},
-    {file = "ruff-0.12.2-py3-none-win32.whl", hash = "sha256:369ffb69b70cd55b6c3fc453b9492d98aed98062db9fec828cdfd069555f5f1a"},
-    {file = "ruff-0.12.2-py3-none-win_amd64.whl", hash = "sha256:dca8a3b6d6dc9810ed8f328d406516bf4d660c00caeaef36eb831cf4871b0639"},
-    {file = "ruff-0.12.2-py3-none-win_arm64.whl", hash = "sha256:48d6c6bfb4761df68bc05ae630e24f506755e702d4fb08f08460be778c7ccb12"},
-    {file = "ruff-0.12.2.tar.gz", hash = "sha256:d7b4f55cd6f325cb7621244f19c873c565a08aff5a4ba9c69aa7355f3f7afd3e"},
+    {file = "ruff-0.12.4-py3-none-linux_armv6l.whl", hash = "sha256:cb0d261dac457ab939aeb247e804125a5d521b21adf27e721895b0d3f83a0d0a"},
+    {file = "ruff-0.12.4-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:55c0f4ca9769408d9b9bac530c30d3e66490bd2beb2d3dae3e4128a1f05c7442"},
+    {file = "ruff-0.12.4-py3-none-macosx_11_0_arm64.whl", hash = "sha256:a8224cc3722c9ad9044da7f89c4c1ec452aef2cfe3904365025dd2f51daeae0e"},
+    {file = "ruff-0.12.4-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e9949d01d64fa3672449a51ddb5d7548b33e130240ad418884ee6efa7a229586"},
+    {file = "ruff-0.12.4-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:be0593c69df9ad1465e8a2d10e3defd111fdb62dcd5be23ae2c06da77e8fcffb"},
+    {file = "ruff-0.12.4-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a7dea966bcb55d4ecc4cc3270bccb6f87a337326c9dcd3c07d5b97000dbff41c"},
+    {file = "ruff-0.12.4-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:afcfa3ab5ab5dd0e1c39bf286d829e042a15e966b3726eea79528e2e24d8371a"},
+    {file = "ruff-0.12.4-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c057ce464b1413c926cdb203a0f858cd52f3e73dcb3270a3318d1630f6395bb3"},
+    {file = "ruff-0.12.4-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e64b90d1122dc2713330350626b10d60818930819623abbb56535c6466cce045"},
+    {file = "ruff-0.12.4-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2abc48f3d9667fdc74022380b5c745873499ff827393a636f7a59da1515e7c57"},
+    {file = "ruff-0.12.4-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:2b2449dc0c138d877d629bea151bee8c0ae3b8e9c43f5fcaafcd0c0d0726b184"},
+    {file = "ruff-0.12.4-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:56e45bb11f625db55f9b70477062e6a1a04d53628eda7784dce6e0f55fd549eb"},
+    {file = "ruff-0.12.4-py3-none-musllinux_1_2_i686.whl", hash = "sha256:478fccdb82ca148a98a9ff43658944f7ab5ec41c3c49d77cd99d44da019371a1"},
+    {file = "ruff-0.12.4-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:0fc426bec2e4e5f4c4f182b9d2ce6a75c85ba9bcdbe5c6f2a74fcb8df437df4b"},
+    {file = "ruff-0.12.4-py3-none-win32.whl", hash = "sha256:4de27977827893cdfb1211d42d84bc180fceb7b72471104671c59be37041cf93"},
+    {file = "ruff-0.12.4-py3-none-win_amd64.whl", hash = "sha256:fe0b9e9eb23736b453143d72d2ceca5db323963330d5b7859d60d101147d461a"},
+    {file = "ruff-0.12.4-py3-none-win_arm64.whl", hash = "sha256:0618ec4442a83ab545e5b71202a5c0ed7791e8471435b94e655b570a5031a98e"},
+    {file = "ruff-0.12.4.tar.gz", hash = "sha256:13efa16df6c6eeb7d0f091abae50f58e9522f3843edb40d56ad52a5a4a4b6873"},
 ]
 
 [[package]]
diff --git a/pyproject.toml b/pyproject.toml
index 3323f4dc..96b99e92 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,14 +11,14 @@ repository = "https://github.com/kozistr/pytorch_optimizer"
 documentation = "https://pytorch-optimizers.readthedocs.io/en/latest"
 keywords = [
     "pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "Alice", "ASGD", "AccSGD", "AdaBelief",
-    "AdaBound", "AdaDelta", "AdaFactor", "AdaGC", "AdaMax", "AdamG", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth",
-    "AdEMAMix", "Simplified-AdEMAMix", "ADOPT", "AdaHessian", "Adai", "Adalite", "AdaLomo", "AdamMini", "AdamP",
-    "AdamS", "Adan", "AggMo", "Aida", "AliG", "Amos", "Apollo", "APOLLO", "AvaGrad", "bSAM", "CAME", "DAdaptAdaGrad",
-    "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DeMo", "DiffGrad", "EXAdam", "FAdam", "Fira", "FOCUS",
-    "Fromage", "FTRL", "GaLore", "Grams", "Gravity", "GrokFast", "GSAM", "Kate", "Lamb", "LaProp", "LARS", "Lion",
-    "LOMO", "Lookahead", "MADGRAD", "MARS", "MSVAG", "Muno", "Nero", "NovoGrad", "OrthoGrad", "PAdam", "PCGrad", "PID",
-    "PNM", "Prodigy", "PSGD", "QHAdam", "QHM", "RACS", "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM", "GCSAM",
-    "LookSAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SCION", "SGDP", "Shampoo",
+    "AdaBound", "AdaDelta", "AdaFactor", "AdaGC", "AdaMax", "AdaMuon", "AdamG", "AdaMod", "AdaNorm", "AdaPNM",
+    "AdaSmooth", "AdEMAMix", "Simplified-AdEMAMix", "ADOPT", "AdaHessian", "Adai", "Adalite", "AdaLomo", "AdamMini",
+    "AdamP", "AdamS", "Adan", "AggMo", "Aida", "AliG", "Amos", "Apollo", "APOLLO", "AvaGrad", "bSAM", "CAME",
+    "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DeMo", "DiffGrad", "EXAdam", "FAdam",
+    "Fira", "FOCUS", "Fromage", "FTRL", "GaLore", "Grams", "Gravity", "GrokFast", "GSAM", "Kate", "Lamb", "LaProp",
+    "LARS", "Lion", "LOMO", "Lookahead", "MADGRAD", "MARS", "MSVAG", "Muno", "Nero", "NovoGrad", "OrthoGrad", "PAdam",
+    "PCGrad", "PID", "PNM", "Prodigy", "PSGD", "QHAdam", "QHM", "RACS", "RAdam", "Ranger", "Ranger21", "RotoGrad",
+    "SAM", "GCSAM", "LookSAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SCION", "SGDP", "Shampoo",
     "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH", "SPAM", "StableSPAM", "SRMM", "StableAdamW",
     "SWATS", "TAM", "Tiger", "TRAC", "VSGD", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1",
     "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD",
diff --git a/pytorch_optimizer/__init__.py b/pytorch_optimizer/__init__.py
index 3f3cd01d..8facc660 100644
--- a/pytorch_optimizer/__init__.py
+++ b/pytorch_optimizer/__init__.py
@@ -89,6 +89,7 @@ from pytorch_optimizer.optimizer import (
     AdaMod,
     AdamP,
     AdamS,
+    AdaMuon,
     AdamW,
     AdamWSN,
     Adan,
diff --git a/pytorch_optimizer/optimizer/__init__.py b/pytorch_optimizer/optimizer/__init__.py
index 20a67cd5..d3eefd40 100644
--- a/pytorch_optimizer/optimizer/__init__.py
+++ b/pytorch_optimizer/optimizer/__init__.py
@@ -66,7 +66,7 @@ from pytorch_optimizer.optimizer.lookahead import Lookahead
 from pytorch_optimizer.optimizer.madgrad import MADGRAD
 from pytorch_optimizer.optimizer.mars import MARS
 from pytorch_optimizer.optimizer.msvag import MSVAG
-from pytorch_optimizer.optimizer.muon import Muon
+from pytorch_optimizer.optimizer.muon import AdaMuon, Muon
 from pytorch_optimizer.optimizer.nero import Nero
 from pytorch_optimizer.optimizer.novograd import NovoGrad
 from pytorch_optimizer.optimizer.orthograd import OrthoGrad
@@ -322,6 +322,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     RACS,
     Alice,
     VSGD,
+    AdaMuon,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 
diff --git a/pytorch_optimizer/optimizer/muon.py b/pytorch_optimizer/optimizer/muon.py
index 8ee2e95b..f44c4939 100644
--- a/pytorch_optimizer/optimizer/muon.py
+++ b/pytorch_optimizer/optimizer/muon.py
@@ -250,3 +250,259 @@ class Muon(BaseOptimizer):
                 p.add_(update, alpha=-step_size)
 
         return loss
+
+
+class AdaMuon(BaseOptimizer):
+    r"""Adaptive Muon Optimizer.
+
+    Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-processing step, in which
+    each 2D parameter's update is replaced with the nearest orthogonal matrix. To efficiently orthogonalize each
+    update, we use a Newton-Schulz iteration, which has the advantage that it can be stably run in bfloat16 on the GPU.
+
+    Muon is intended to optimize only the internal ≥2D parameters of a network. Embeddings, classifier heads, and
+    scalar or vector parameters should be optimized using AdamW.
+
+    Some warnings:
+    - We believe this optimizer is unlikely to work well for training with small batch size.
+    - We believe it may not work well for fine-tuning pretrained models, but we haven't tested this.
+
+    :param params: PARAMETERS. the parameters to be optimized by Muon.
+    :param lr: float. learning rate.
+    :param betas: The betas for the Muon.
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    :param nesterov: bool. whether to use nesterov momentum.
+    :param ns_steps: int. the number of Newton-Schulz iterations to run. (5 is probably always enough)
+    :param use_adjusted_lr: bool. whether to use adjusted learning rate, which is from the Moonlight.
+        reference: https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py
+    :param adamw_params: Optional[PARAMETERS] The parameters to be optimized by AdamW. Any parameters in `muon_params`
+        which are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. It'd be
+        better to create AdamW optimizer instead of using this.
+    :param adamw_betas: The betas for the internal AdamW.
+    :param adamw_lr: float. The learning rate for the internal AdamW.
+    :param adamw_wd: float. The weight decay for the internal AdamW.
+    :param adamw_eps: float. The epsilon for the internal AdamW.
+    :param maximize: bool. maximize the objective with respect to the params, instead of minimizing.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        betas: BETAS = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+        weight_decouple: bool = True,
+        nesterov: bool = True,
+        ns_steps: int = 5,
+        use_adjusted_lr: bool = False,
+        adamw_params: Optional[PARAMETERS] = None,
+        adamw_betas: BETAS = (0.9, 0.999),
+        adamw_lr: float = 3e-4,
+        adamw_wd: float = 0.0,
+        eps: float = 1e-8,
+        maximize: bool = False,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_learning_rate(adamw_lr)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_positive(ns_steps, 'ns_steps')
+        self.validate_betas(betas)
+        self.validate_betas(adamw_betas)
+        self.validate_non_negative(adamw_wd, 'adamw_wd')
+        self.validate_non_negative(eps, 'eps')
+
+        params = self.get_parameters(params)
+        adamw_params = self.get_parameters(adamw_params) if adamw_params is not None else []
+        params.extend(adamw_params)
+
+        self.world_size: int = int(os.environ.get('WORLD_SIZE', '1'))
+        self.rank: int = int(os.environ.get('RANK', '0'))
+        self.maximize = maximize
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'nesterov': nesterov,
+            'ns_steps': ns_steps,
+            'use_adjusted_lr': use_adjusted_lr,
+            'adamw_lr': adamw_lr,
+            'adamw_lr_ratio': adamw_lr / lr,
+            'adamw_betas': adamw_betas,
+            'adamw_wd': adamw_wd,
+            'eps': eps,
+        }
+
+        super().__init__(params, defaults)
+
+        self.set_muon_state(params, adamw_params)
+
+    def __str__(self) -> str:
+        return 'AdaMuon'
+
+    @staticmethod
+    def get_parameters(params: PARAMETERS) -> List[torch.Tensor]:
+        if isinstance(params, list) and isinstance(params[0], torch.Tensor):
+            return params
+
+        new_params = []
+        for group in params:
+            if isinstance(group, dict) and 'params' in group:
+                new_params.extend(list(group['params']))
+            else:
+                new_params.append(group)
+
+        return new_params
+
+    def set_muon_state(self, params: PARAMETERS, adamw_params: PARAMETERS) -> None:
+        r"""Set use_muon flag."""
+        for p in params:
+            self.state[p]['use_muon'] = p.ndim >= 2
+
+        for p in adamw_params:
+            self.state[p]['use_muon'] = False
+
+    def init_group(self, group: GROUP, **kwargs) -> None:
+        pass
+
+    @staticmethod
+    def get_adjusted_lr(lr: float, param_shape: Tuple[float, ...], use_adjusted_lr: bool = False) -> float:
+        r"""Get the adjust learning rate."""
+        output_shape, *input_shape = param_shape
+        input_shape = math.prod(input_shape)
+
+        ratio: float = (
+            math.pow(max(1.0, output_shape / input_shape), 0.5)
+            if use_adjusted_lr
+            else 0.2 * math.sqrt(max(output_shape, input_shape))
+        )
+
+        return lr * ratio
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if 'step' not in group:
+                self.init_group(group)
+                group['step'] = 1
+            else:
+                group['step'] += 1
+
+            params = []
+            for p in group['params']:
+                if p.grad is not None and self.state[p]['use_muon']:
+                    if p.grad.is_sparse:
+                        raise NoSparseGradientError(str(self))
+                    if torch.is_complex(p):
+                        raise NoComplexParameterError(str(self))
+                    params.append(p)
+
+            if len(params) == 0:
+                continue
+
+            beta1, beta2 = group['betas']
+
+            bias_correction2: float = 1.0 - self.debias(beta2, group['step'])
+
+            total_params: int = sum(p.numel() for p in params)
+            updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16)
+            curr_idx: int = 0
+
+            for i, p in enumerate(params):
+                if i % self.world_size != self.rank:
+                    curr_idx += p.numel()
+                    continue
+
+                grad = p.grad
+                if grad.ndim > 2:
+                    grad = grad.view(grad.size(0), -1)
+
+                state = self.state[p]
+                if 'm' not in state:
+                    state['m'] = torch.zeros_like(grad)
+                    state['v'] = torch.zeros_like(grad.flatten())
+
+                buf = state['m']
+                buf.lerp_(grad, weight=1.0 - beta1)
+
+                grad = grad.lerp_(buf, beta1) if group['nesterov'] else buf
+
+                grad = zero_power_via_newton_schulz_5(grad, num_steps=group['ns_steps']).flatten()
+
+                updates_flat[curr_idx:curr_idx + p.numel()] = grad  # fmt: skip
+
+            if self.world_size > 1:  # pragma: no cover
+                all_reduce(updates_flat, op=ReduceOp.SUM)
+
+            curr_idx: int = 0
+            for p in params:
+                if 'v' not in self.state[p]:
+                    continue
+
+                g = updates_flat[curr_idx:curr_idx + p.numel()]  # fmt: skip
+
+                v = self.state[p]['v']
+                v.mul_(beta2).addcmul_(g, g, value=1.0 - beta2)
+
+                update = g / (v / bias_correction2).sqrt().add_(group['eps'])
+                update = update.view_as(p)
+
+                update.mul_(0.2 * math.sqrt(p.numel())).div_(update.norm().add_(group['eps']))
+
+                lr: float = self.get_adjusted_lr(group['lr'], p.size(), group['use_adjusted_lr'])
+
+                self.apply_weight_decay(
+                    p,
+                    grad=g.view_as(p),
+                    lr=lr,
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=False,
+                )
+
+                p.add_(update, alpha=-lr)
+                curr_idx += p.numel()
+
+            params = [p for p in group['params'] if p.grad is not None and not self.state[p]['use_muon']]
+
+            lr: float = group['adamw_lr_ratio'] * group['lr']
+            beta1, beta2 = group['adamw_betas']
+
+            bias_correction1: float = self.debias(beta1, group['step'])
+            bias_correction2: float = self.debias(beta2, group['step'])
+
+            scale: float = bias_correction1 / bias_correction2 ** 0.5  # fmt: skip
+            step_size: float = lr / scale
+
+            for p in params:
+                grad = p.grad
+                state = self.state[p]
+                if 'exp_avg' not in state:
+                    state['exp_avg'] = torch.zeros_like(grad)
+                    state['exp_avg_sq'] = torch.zeros_like(grad)
+
+                buf1, buf2 = state['exp_avg'], state['exp_avg_sq']
+                buf1.lerp_(grad, weight=1.0 - beta1)
+                buf2.lerp_(grad.square(), weight=1.0 - beta2)
+
+                update = buf1 / buf2.sqrt().add_(group['eps'])
+
+                self.apply_weight_decay(
+                    p,
+                    grad,
+                    lr=lr,
+                    weight_decay=group['adamw_wd'],
+                    weight_decouple=True,
+                    fixed_decay=False,
+                )
+
+                p.add_(update, alpha=-step_size)
+
+        return loss


Test patch:
diff --git a/tests/constants.py b/tests/constants.py
index ae62691b..5361f1de 100644
--- a/tests/constants.py
+++ b/tests/constants.py
@@ -42,6 +42,7 @@ from pytorch_optimizer.optimizer import (
     AdaMod,
     AdamP,
     AdamS,
+    AdaMuon,
     AdamWSN,
     Adan,
     AdaNorm,
@@ -234,6 +235,7 @@ BETA_OPTIMIZER_NAMES: List[str] = [
     'ademamix',
     'soap',
     'muon',
+    'adamuon',
     'laprop',
     'apollo',
     'mars',
@@ -598,6 +600,8 @@ OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (FTRL, {'lr': 1e0, 'beta': 0.0, 'lambda_1': 0.0, 'lambda_2': 0.0}, 5),
     (Muon, {'lr': 5e0, 'use_adjusted_lr': True, 'adam_lr': 1e0, 'adamw_wd': 1e-2}, 5),
     (Muon, {'lr': 2e0, 'adam_lr': 1e0, 'adamw_wd': 1e-2, 'nesterov': False}, 5),
+    (AdaMuon, {'lr': 5e0, 'use_adjusted_lr': True, 'adam_lr': 1e0, 'adamw_wd': 1e-2}, 5),
+    (AdaMuon, {'lr': 5e0, 'adam_lr': 1e0, 'adamw_wd': 1e-2, 'nesterov': False}, 7),
     (LaProp, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (LaProp, {'lr': 1e0, 'centered': True, 'weight_decay': 1e-3}, 11),
     (LaProp, {'lr': 1e0, 'ams_bound': True, 'weight_decay': 1e-3}, 5),
diff --git a/tests/test_load_modules.py b/tests/test_load_modules.py
index 5f86256e..7827589d 100644
--- a/tests/test_load_modules.py
+++ b/tests/test_load_modules.py
@@ -34,9 +34,9 @@ def test_load_lr_scheduler_invalid(invalid_lr_scheduler_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 106
-    assert len(get_supported_optimizers('adam*')) == 10
-    assert len(get_supported_optimizers(['adam*', 'ranger*'])) == 13
+    assert len(get_supported_optimizers()) == 107
+    assert len(get_supported_optimizers('adam*')) == 11
+    assert len(get_supported_optimizers(['adam*', 'ranger*'])) == 14
 
 
 def test_get_supported_lr_schedulers():
diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py
index 959f1e8d..d5e385b3 100644
--- a/tests/test_optimizers.py
+++ b/tests/test_optimizers.py
@@ -6,7 +6,7 @@ import torch
 from torch import nn
 
 from pytorch_optimizer.base.exception import NoClosureError, ZeroParameterSizeError
-from pytorch_optimizer.optimizer import DynamicLossScaler, Muon, load_optimizer
+from pytorch_optimizer.optimizer import DynamicLossScaler, load_optimizer
 from pytorch_optimizer.optimizer.alig import l2_projection
 from pytorch_optimizer.optimizer.grokfast import gradfilter_ema, gradfilter_ma
 from pytorch_optimizer.optimizer.scion import build_lmo_norm
@@ -29,7 +29,7 @@ from tests.utils import (
 def build_optimizer_parameter(parameters, optimizer_name, config):
     if optimizer_name == 'AliG':
         config.update({'projection_fn': lambda: l2_projection(parameters, max_norm=1)})
-    if optimizer_name == 'Muon':
+    if optimizer_name in ('Muon', 'AdaMuon'):
         adamw_params = [p for i, p in enumerate(parameters) if i >= 2]
         parameters = [p for i, p in enumerate(parameters) if i < 2]
         config.update({'adamw_params': adamw_params})
@@ -474,8 +474,9 @@ def test_soap_merge_dims_channel_last(environment):
         optimizer.step()
 
 
+@pytest.mark.parametrize('optimizer_name', ['Muon', 'AdaMuon'])
 @pytest.mark.parametrize('rank', ['1', '0'])
-def test_muon_rank(rank):
+def test_muon_rank(optimizer_name, rank):
     os.environ['RANK'] = rank
 
     model = nn.Sequential(
@@ -484,7 +485,7 @@ def test_muon_rank(rank):
         nn.Conv2d(1, 1, (2, 2)),
     )
 
-    optimizer = Muon(model.parameters())
+    optimizer = load_optimizer(optimizer_name)(model.parameters())
     optimizer.zero_grad()
 
     model[0].weight.grad = torch.randn(1, 1, 1)