From bc4b7ede842a3bf2d656dea5141c91c8589a47d6 Mon Sep 17 00:00:00 2001 From: Dmitry Voropaev Date: Tue, 20 Oct 2026 04:56:38 +0300 Subject: [PATCH] Fix AdaptiveMuon step for param groups added after init AdaptiveMuon set beta2 and eps with setdefault on the param groups that existed when the constructor ran, but did not put them in the optimizer defaults. torch.optim.Optimizer.add_param_group fills missing hyperparameters from self.defaults only, so a group added later had no beta2 or eps and step() failed with KeyError: 'beta2'. Pass beta2 and eps to the base class as keyword arguments so they land in the defaults like lr and momentum, and drop the setdefault loop. Groups that set their own beta2 or eps keep them. Signed-off-by: Dmitry Voropaev --- .../adaptive_muon.py | 6 ++--- tests/test_adaptive_muon.py | 26 +++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py b/emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py index 892c625f..4db90753 100644 --- a/emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py +++ b/emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py @@ -111,13 +111,11 @@ def raw_orthogonalize_fn(grad: torch.Tensor) -> torch.Tensor: weight_decay_method=weight_decay_method, fp32_matmul_prec=fp32_matmul_prec, scaled_orthogonalize_fn=raw_orthogonalize_fn, + beta2=beta2, + eps=eps, ) self.moment2_method = moment2_method - for group in self.param_groups: - group.setdefault("beta2", beta2) - group.setdefault("eps", eps) - def _apply_muon_scale(self, update: torch.Tensor, size_out: int, size_in: int) -> torch.Tensor: scale_factor = muon.get_muon_scale_factor(size_out, size_in, mode=self.scale_mode) logging.debug("Applying Muon scale factor %s, extra_scale_factor=%s", scale_factor, self.extra_scale_factor) diff --git a/tests/test_adaptive_muon.py b/tests/test_adaptive_muon.py index 5d97e4f3..70044449 100644 --- a/tests/test_adaptive_muon.py +++ b/tests/test_adaptive_muon.py @@ -103,6 +103,32 @@ def test_moment2_matches_shapes(self, shape, moment2_method) -> None: elif moment2_method == "namo": self.assertEqual(moment2.shape, torch.Size([1])) + @parameterized.parameters( + *({"moment2_method": moment2_method} for moment2_method in MOMENT2_METHODS), + ) + def test_step_with_param_group_added_after_init(self, moment2_method) -> None: + first_param = nn.Parameter(torch.randn(8, 16, dtype=torch.float32, device=FLAGS.device)) + first_param.grad = torch.randn_like(first_param) + adaptive_opt = AdaptiveMuon( + [first_param], + lr=0.01, + momentum=0.9, + weight_decay=0.0, + moment2_method=moment2_method, + beta2=0.9, + eps=1e-6, + ) + + added_param = nn.Parameter(torch.randn(16, 8, dtype=torch.float32, device=FLAGS.device)) + added_param.grad = torch.randn_like(added_param) + adaptive_opt.add_param_group({"params": [added_param]}) + + adaptive_opt.step() + + self.assertEqual(adaptive_opt.param_groups[1]["beta2"], 0.9) + self.assertEqual(adaptive_opt.param_groups[1]["eps"], 1e-6) + self.assertIn("moment2_buffer", adaptive_opt.state[added_param]) + @parameterized.parameters( *({"moment2_method": moment2_method} for moment2_method in MOMENT2_METHODS), )