From 5aabbb199181531d60095ad2e810dc154b3fbd79 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Mon, 24 Aug 2026 13:33:10 -0700
Subject: [PATCH] fix Muon weight decay being a no-op (#17709)

* fix Muon weight decay being a no-op

LARS._step computed the post-momentum weight decayed param but only
used it for a dtype cast, so the decay was never applied. Fold the
decay into the update instead, matching torch's param.mul_(1 - lr*wd).

test_muon_wd passed anyway since lr*wd=1e-5 is far below atol, so
also bump the test's weight_decay to 10 to actually exercise it.

* muon: apply weight decay after lr scaling

keeps the decoupled decay independent of the LARS trust ratio r,
matching torch's param.mul_(1 - lr*wd). no behavior change today since
r is always 1.0 on the pre_wd=False path (Muon has tcoef=0).
---
 test/backend/test_optim.py | 3 ++-
 tinygrad/nn/optim.py       | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/test/backend/test_optim.py b/test/backend/test_optim.py
index 299ec7a44d..74ce1d1f25 100644
--- a/test/backend/test_optim.py
+++ b/test/backend/test_optim.py
@@ -87,7 +87,8 @@ class TestOptim(unittest.TestCase):
   def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
   # TODO: disabled due to big atol
   # def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
-  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
+  # NOTE: big weight_decay so a missing wd would be way over atol
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 10}, 1e-3, 3e-4)
   # TODO: disabled due to big atol
   # def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)
 
diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py
index b07fd75f13..2eca8b3e6e 100644
--- a/tinygrad/nn/optim.py
+++ b/tinygrad/nn/optim.py
@@ -121,10 +121,10 @@ class LARS(Optimizer):
         self.b[i].assign(self.momentum * self.b[i] + g)  # NOTE: self.b[i] is zero on the first run, no if required
         g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i]
       if self.ns_coefficients: g = g.reshape(g.shape[0], -1).newton_schulz(self.ns_steps, self.ns_coefficients).reshape(g.shape)
-      # muon does post momentum weight decay
-      if not self.pre_wd and self.wd > 0: t = t.detach() * (1.0 - self.wd * self.lr)
       # popular momentum does pre learning rate update
       if not self.classic: g = g * r * self.lr
+      # muon does post momentum weight decay
+      if not self.pre_wd and self.wd > 0: g = g + self.wd * self.lr * t.detach()
       ret.append(g.cast(t.dtype))
     return ret, self.b