Improve consistency among couple of orthogonalized optimizers (#79)

Signed-off-by: Hao Wu <[email protected]>

Author: skyw

emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py

@@ -195,7 +195,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
         for group in self.param_groups:
             for p in group["params"]:
                 if p.dim() != 2:
-                    raise ValueError("AdaptiveMuon only supports 2D parameters")
+                    raise ValueError(f"{self.__class__.__name__} only supports 2D parameters")
                 grad = p.grad
                 if grad is None:
                     continue
@@ -223,7 +223,8 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                     grad = exp_avg
 
                 with utils.fp32_matmul_precision(self.fp32_matmul_prec):
-                    orth_grad = self.scaled_orthogonalize_fn(grad)
+                    group_kwargs = {k: v for k, v in group.items() if k != "params"}
+                    orth_grad = self.orthogonalize(p, grad, **group_kwargs)
 
                 update = self._apply_moment2_normalization(
                     orth_grad=orth_grad,

emerging_optimizers/orthogonalized_optimizers/orthogonalized_optimizer.py

@@ -140,8 +140,8 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
 
         for group in self.param_groups:
             for p in group["params"]:
-                if p.dim() == 1:
-                    raise ValueError(f"{self.__class__.__name__} does not support 1D parameters")
+                if p.dim() != 2:
+                    raise ValueError(f"{self.__class__.__name__} only supports 2D parameters")
                 grad = p.grad
                 if grad is None:
                     continue
@@ -172,11 +172,11 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
 
                 with utils.fp32_matmul_precision(self.fp32_matmul_prec):
                     group_kwargs = {k: v for k, v in group.items() if k != "params"}
-                    grad = self.orthogonalize(p, grad, **group_kwargs)
+                    orth_grad = self.orthogonalize(p, grad, **group_kwargs)
 
                 # perform weight update
                 # scale is applied to have update RMS == 1
-                p.add_(grad, alpha=-group["lr"])
+                p.add_(orth_grad, alpha=-group["lr"])
 
         return loss