Merge pull request #373 from kozistr/fix/muon-optimizer

[Fix] Correct the learning rate ratio in `Muon` optimizer

Author: kozistr

docs/changelogs/v3.5.1.md

@@ -7,3 +7,7 @@
 ### Update
 
 * Update `SCION` optimizer based on the official implementation. (#369)
+
+### Fix
+
+* Correct the learning rate ratio in `Muon` optimizer properly. (#371, #372, #373)

poetry.lock

@@ -1,6 +1,6 @@
 import math
 import os
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import torch
 from torch.distributed import ReduceOp, all_reduce
@@ -131,9 +131,18 @@ def reset(self):
                 state['moment2'] = torch.zeros_like(p)
 
     @staticmethod
-    def adjust_lr_for_muon(lr: float, param_shape) -> float:
-        adjusted_ratio: float = 0.2 * math.sqrt(max(param_shape[0], param_shape[1]))
-        return lr * adjusted_ratio
+    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:
@@ -202,9 +211,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     fixed_decay=False,
                 )
 
-                lr: float = self.adjust_lr_for_muon(group['lr'], p.size()) if group['use_adjusted_lr'] else group['lr']
+                lr: float = self.get_adjusted_lr(group['lr'], p.size(), group['use_adjusted_lr'])
 
-                p.add_(g, alpha=-lr * (max(1.0, p.size(-2) / p.size(-1)) ** 0.5))
+                p.add_(g, 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']]

pytorch_optimizer/optimizer/muon.py

@@ -5,12 +5,13 @@ black==25.1.0 ; python_version >= "3.9"
 click==8.1.8 ; python_version >= "3.8"
 colorama==0.4.6 ; python_version >= "3.8" and (sys_platform == "win32" or platform_system == "Windows")
 coverage[toml]==7.6.1 ; python_version == "3.8"
-coverage[toml]==7.6.12 ; python_version >= "3.9"
+coverage[toml]==7.8.0 ; python_version >= "3.9"
 exceptiongroup==1.2.2 ; python_version < "3.11" and python_version >= "3.8"
 filelock==3.16.1 ; python_version == "3.8"
 filelock==3.18.0 ; python_version >= "3.9"
-fsspec==2025.3.0 ; python_version >= "3.8"
-iniconfig==2.0.0 ; python_version >= "3.8"
+fsspec==2025.3.0 ; python_version == "3.8"
+fsspec==2025.3.2 ; python_version >= "3.9"
+iniconfig==2.1.0 ; python_version >= "3.8"
 isort==5.13.2 ; python_version == "3.8"
 isort==6.0.1 ; python_version >= "3.9"
 jinja2==3.1.6 ; python_version >= "3.8"
@@ -22,17 +23,18 @@ networkx==3.1 ; python_version == "3.8"
 networkx==3.2.1 ; python_version >= "3.9"
 numpy==1.24.4 ; python_version == "3.8"
 numpy==2.0.2 ; python_version >= "3.9"
-packaging==24.2 ; python_version >= "3.8"
+packaging==25.0 ; python_version >= "3.8"
 pathspec==0.12.1 ; python_version >= "3.8"
-platformdirs==4.3.6 ; python_version >= "3.8"
+platformdirs==4.3.6 ; python_version == "3.8"
+platformdirs==4.3.7 ; python_version >= "3.9"
 pluggy==1.5.0 ; python_version >= "3.8"
 pytest-cov==5.0.0 ; python_version >= "3.8"
 pytest==8.3.5 ; python_version >= "3.8"
-ruff==0.11.0 ; python_version >= "3.8"
-setuptools==76.0.0 ; python_version >= "3.12"
+ruff==0.11.6 ; python_version >= "3.8"
+setuptools==79.0.0 ; python_version >= "3.12"
 sympy==1.13.1 ; python_version >= "3.9"
 sympy==1.13.3 ; python_version == "3.8"
 tomli==2.2.1 ; python_full_version <= "3.11.0a6" and python_version >= "3.8"
 torch==2.4.1+cpu ; python_version == "3.8"
 torch==2.6.0+cpu ; python_version >= "3.9"
-typing-extensions==4.12.2 ; python_version >= "3.8"
+typing-extensions==4.13.2 ; python_version >= "3.8"

requirements-dev.txt

@@ -2,7 +2,8 @@
 
 filelock==3.16.1 ; python_version == "3.8"
 filelock==3.18.0 ; python_version >= "3.9"
-fsspec==2025.3.0 ; python_version >= "3.8"
+fsspec==2025.3.0 ; python_version == "3.8"
+fsspec==2025.3.2 ; python_version >= "3.9"
 jinja2==3.1.6 ; python_version >= "3.8"
 markupsafe==2.1.5 ; python_version == "3.8"
 markupsafe==3.0.2 ; python_version >= "3.9"
@@ -11,9 +12,9 @@ networkx==3.1 ; python_version == "3.8"
 networkx==3.2.1 ; python_version >= "3.9"
 numpy==1.24.4 ; python_version == "3.8"
 numpy==2.0.2 ; python_version >= "3.9"
-setuptools==76.0.0 ; python_version >= "3.12"
+setuptools==79.0.0 ; python_version >= "3.12"
 sympy==1.13.1 ; python_version >= "3.9"
 sympy==1.13.3 ; python_version == "3.8"
 torch==2.4.1+cpu ; python_version == "3.8"
 torch==2.6.0+cpu ; python_version >= "3.9"
-typing-extensions==4.12.2 ; python_version >= "3.8"
+typing-extensions==4.13.2 ; python_version >= "3.8"

requirements.txt

@@ -525,7 +525,7 @@
     (ADOPT, {'lr': 1e0}, 5),
     (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': 1e0, 'adam_lr': 1e0, 'adamw_wd': 1e-2, 'nesterov': False}, 5),
+    (Muon, {'lr': 2e0, 'adam_lr': 1e0, 'adamw_wd': 1e-2, 'nesterov': False}, 5),
     (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),