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