Fix AdaptiveMuon step for param groups added after init

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)

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