Merge branch 'main' into skyw/remove-fstring-in-log

Author: skyw

emerging_optimizers/orthogonalized_optimizers/muon_utils.py

@@ -215,7 +215,9 @@ def newton_schulz(
         if not normalize_in_double:
             X = torch.nn.functional.normalize(x, p=2, dim=(-2, -1), eps=eps)  # type: ignore[arg-type]
         else:
-            # eps is ignored when normalize in double.
+            # eps is ignored when normalize in double so that zero division can happen if norm is exact 0.
+            # However, if norm is 0 in double precision, it means the entire input is 0, which usually
+            # suggests something wrong has happened in training. So we don't guard it here.
             norm = torch.linalg.vector_norm(x, dim=(-2, -1), keepdim=True, dtype=torch.float64).to(x.dtype)
             X = x / norm
 

tests/soap_reference.py

@@ -278,8 +278,6 @@ def update_preconditioner(self, grad, state, max_precond_dim=10000, merge_dims=F
                 state["exp_avg"], state, merge_dims=merge_dims, max_precond_dim=max_precond_dim
             )
 
-        # print("wtf1", state["exp_avg"])
-
     def project_back(self, grad, state, merge_dims=False, max_precond_dim=10000):
         """
         Projects the gradient back to the original space.