forked from tinygrad/tinygrad
George HotzandGitHub
5aabbb1991
* 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).
Three groups of tests run in CI backend -- tests that run on each backend null -- tests that don't require any backend unit -- tests that only run on a single backend in CI