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

@@ -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)
 

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