remove svd exp

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

emerging_optimizers/soap/rekls.py

@@ -76,8 +76,8 @@ class TpRekls(opt_mixin.WeightDecayMixin, optim.Optimizer):
 
     Reimplemented from scratch (not inheriting from :class:`~emerging_optimizers.soap.soap.SOAP`) so the
     tensor-parallel bookkeeping stays isolated. Eigenbases are not stored in optimizer state; they are
-    recomputed via :func:`~emerging_optimizers.soap.soap_utils.get_eigenbasis_svd` from the kronecker
-    factors. Each step calls svd twice — once on the pre-update L, R for the
+    recomputed via :func:`~emerging_optimizers.soap.soap_utils.get_eigenbasis_eigh` from the kronecker
+    factors. Each step calls eigh twice — once on the pre-update L, R for the
     :func:`~emerging_optimizers.soap.soap.update_kronecker_factors_kl_shampoo` correction, and once on
     the post-update L, R for the gradient projection.
 
@@ -231,9 +231,9 @@ def step(self, closure: None = None) -> None:
                 shampoo_beta = 1 - (1 - shampoo_beta) / (1 - shampoo_beta**curr_iter_1_based)
 
                 # KL-Shampoo correction needs the eigenbasis of the *pre-update* L, R; recompute it
-                # via svd since we do not persist eigenbases across steps.
+                # via eigh since we do not persist eigenbases across steps.
                 with utils.fp32_matmul_precision(self.fp32_matmul_prec):
-                    pre_eigenbasis_list = soap_utils.get_eigenbasis_svd(kronecker_factor_list)
+                    pre_eigenbasis_list = soap_utils.get_eigenbasis_eigh(kronecker_factor_list)
                     soap.update_kronecker_factors_kl_shampoo(
                         kronecker_factor_list,
                         full_grad,
@@ -249,11 +249,15 @@ def step(self, closure: None = None) -> None:
                     state["R"].copy_(kronecker_factor_list[1].chunk(self.tp_size, dim=0)[self.tp_rank])
 
                 with utils.fp32_matmul_precision(self.fp32_matmul_prec):
-                    # Rotate exp_avg from the pre-update eigenbasis to the post-update eigenbasis
-                    # (matches update_eigenbasis_and_exp_avgs in soap.py for the eigh path; we use svd here).
-                    state["exp_avg"] = soap.precondition(state["exp_avg"], pre_eigenbasis_list, dims=[[0], [1]])
-                    eigenbasis_list = soap_utils.get_eigenbasis_svd(kronecker_factor_list)
-                    state["exp_avg"] = soap.precondition(state["exp_avg"], eigenbasis_list, dims=[[0], [0]])
+                    # Rotate exp_avg from the pre-update eigenbasis to the post-update eigenbasis,
+                    # and recompute the post-update eigenbasis via eigh.
+                    eigenbasis_list, state["exp_avg"], state["exp_avg_sq"] = soap.update_eigenbasis_and_exp_avgs(
+                        kronecker_factor_list=kronecker_factor_list,
+                        eigenbasis_list=pre_eigenbasis_list,
+                        exp_avg_sq=state["exp_avg_sq"],
+                        exp_avg=state["exp_avg"],
+                        use_eigh=True,
+                    )
 
                     full_grad_projected = soap.precondition(full_grad, eigenbasis_list, dims=[[0], [0]])