ai fixes. DONOT merge yet

skyw · skyw · commit 27fe09b7be8a · 2026-07-02T19:54:04.000-07:00
Signed-off-by: Hao Wu &lt;skyw@nvidia.com&gt;
diff --git a/emerging_optimizers/soap/moso.py b/emerging_optimizers/soap/moso.py
@@ -237,16 +237,20 @@ def _update_eigenbasis_and_adam_exp_avgs(
     if use_eigh:
         _, (updated_eigenbasis,) = soap_utils.get_eigenbasis_eigh([momentum_factor])
     else:
-        # get_eigenbasis_qr permutes exp_avg_sq along the factor-list axes (dim 0 for the single
-        # factor here), so transpose the right-preconditioned case in and out.
+        # permute_eigenbasis_and_exp_avg_sq permutes exp_avg_sq along the factor-list axes (dim 0 for
+        # the single factor here), so transpose the right-preconditioned case in and out.
         x = exp_avg_sq if left_preconditioned else exp_avg_sq.mT
-        _, (updated_eigenbasis,), x = soap_utils.get_eigenbasis_qr(
+        (eigenbasis,), x = soap_utils.permute_eigenbasis_and_exp_avg_sq(
             [momentum_factor],
             [eigenbasis],
             x,
-            power_iter_steps=power_iter_steps,
         )
         exp_avg_sq = x if left_preconditioned else x.mT
+        _, (updated_eigenbasis,) = soap_utils.get_eigenbasis_qr(
+            [momentum_factor],
+            [eigenbasis],
+            power_iter_steps=power_iter_steps,
+        )
 
     exp_avg = _project_to_one_sided_eigenbasis(
         x=exp_avg,
diff --git a/emerging_optimizers/soap/soap.py b/emerging_optimizers/soap/soap.py
@@ -455,8 +455,9 @@ def update_eigenbasis_and_exp_avgs(
     1. Projects exp_avg back to the original basis
     2. Updates the eigenbases (via eigh, or QR decomposition and power iteration (orthogonal iteration))
        along with the (approximate) eigenvalues of the kronecker factors in the updated eigenbases;
-       on the QR path exp_avg_sq is permuted to match the sorted eigenbasis columns (not needed for
-       the eigh path, which rebuilds the basis from scratch)
+       on the QR path the eigenbases and exp_avg_sq slots are first permuted by descending approximate
+       eigenvalues of the updated factors (not needed for the eigh path, which rebuilds the basis
+       from scratch)
     3. Projects exp_avg back to the new eigenbasis
 
     Args:
@@ -505,10 +506,16 @@ def update_eigenbasis_and_exp_avgs(
             kronecker_factor_list,
         )
     else:
-        updated_eigvals_list, updated_eigenbasis_list, exp_avg_sq = soap_utils.get_eigenbasis_qr(
+        # Orthogonal iteration is column-order sensitive, so first permute the eigenbases (and the
+        # matching exp_avg_sq slots) by descending approximate eigenvalues of the updated factors.
+        eigenbasis_list, exp_avg_sq = soap_utils.permute_eigenbasis_and_exp_avg_sq(
             kronecker_factor_list,
             eigenbasis_list,
             exp_avg_sq,
+        )
+        updated_eigvals_list, updated_eigenbasis_list = soap_utils.get_eigenbasis_qr(
+            kronecker_factor_list,
+            eigenbasis_list,
             power_iter_steps,
         )
 
diff --git a/emerging_optimizers/soap/soap_utils.py b/emerging_optimizers/soap/soap_utils.py
@@ -27,6 +27,7 @@
     "get_eigenbasis_eigh",
     "get_eigenbasis_qr",
     "get_eigenbasis_svd",
+    "permute_eigenbasis_and_exp_avg_sq",
 ]
 
 
@@ -78,47 +79,70 @@ def get_eigenbasis_svd(
     return updated_eigenbasis_list
 
 
-def get_eigenbasis_qr(
+def permute_eigenbasis_and_exp_avg_sq(
     kronecker_factor_list: TensorList,
     eigenbasis_list: TensorList,
     exp_avg_sq: torch.Tensor,
+) -> tuple[TensorList, torch.Tensor]:
+    """Permute each eigenbasis and the matching ``exp_avg_sq`` axis by descending approximate eigenvalues.
+
+    Computes the approximate eigenvalues of each eigenbasis against its kronecker factor (the Rayleigh
+    quotients :math:`\\mathrm{diag}(Q^{\\top} K Q)`) and permutes the eigenbasis columns into descending
+    order, permuting ``exp_avg_sq`` along the corresponding axis so its per-slot statistics stay
+    aligned. Used before :func:`get_eigenbasis_qr`, whose orthogonal iteration is column-order
+    sensitive; the eigh path does not need it since it rebuilds the eigenbasis from scratch.
+
+    Args:
+        kronecker_factor_list: List of preconditioner matrices (L and R).
+        eigenbasis_list: List of current eigenbases (QL and QR).
+        exp_avg_sq: Inner Adam second moment tensor, permuted along each kronecker-factor axis to
+            match the new descending-eigenvalue column ordering.
+
+    Returns:
+        ``(permuted_eigenbasis_list, permuted_exp_avg_sq)``.
+    """
+    permuted_eigenbasis_list: TensorList = []
+    for ind, (kronecker_factor, eigenbasis) in enumerate(zip(kronecker_factor_list, eigenbasis_list, strict=True)):
+        approx_eigvals = eig_utils.conjugate(kronecker_factor, eigenbasis, diag=True)
+        sort_idx = torch.argsort(approx_eigvals, descending=True)
+        permuted_eigenbasis_list.append(eigenbasis[:, sort_idx])
+        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
+    return permuted_eigenbasis_list, exp_avg_sq
+
+
+def get_eigenbasis_qr(
+    kronecker_factor_list: TensorList,
+    eigenbasis_list: TensorList,
     power_iter_steps: int = 1,
-) -> tuple[TensorList, TensorList, torch.Tensor]:
+) -> tuple[TensorList, TensorList]:
     """Updates the eigenbases of the preconditioner using power iteration and QR.
 
-    Each eigenbasis is refined with multiple rounds of power iteration followed by QR decomposition
-    (orthogonal iteration), then its columns are sorted by descending approximate eigenvalues of the
-    kronecker factor (the Rayleigh quotients :math:`\\mathrm{diag}(Q^{\\top} K Q)`), so the returned
-    eigenvalues and eigenbases are in descending order like :func:`get_eigenbasis_eigh`.
-    ``exp_avg_sq`` is permuted along each kronecker-factor axis by the same sort, so its per-slot
-    statistics stay aligned with the eigenbasis columns.
+    Computes using multiple rounds of power iteration followed by QR decomposition (orthogonal
+    iteration). ``eigenbasis_list`` is expected to be already sorted by descending approximate
+    eigenvalues (see :func:`permute_eigenbasis_and_exp_avg_sq`). The returned approximate eigenvalues
+    (the Rayleigh quotients :math:`\\mathrm{diag}(Q^{\\top} K Q)`) are aligned with the updated
+    eigenbasis columns.
 
     Args:
         kronecker_factor_list: List of preconditioner matrices (L and R).
         eigenbasis_list: List of current eigenbases (QL and QR).
-        exp_avg_sq: Inner Adam second moment tensor, permuted along each kronecker-factor axis to
-            match the sorted eigenbasis columns.
         power_iter_steps: Number of power iteration steps to perform before QR decomposition.
             More steps can lead to better convergence but increased computation time.
 
     Returns:
         Tuple of (list of approximate eigenvalues of each kronecker factor in its updated eigenbasis,
-        in descending order, updated list of orthonormal eigenbases (QL and QR) with columns ordered to
-        match, ``exp_avg_sq`` permuted to match).
+        updated list of orthonormal eigenbases (QL and QR)).
     """
     updated_eigenbasis_list: TensorList = []
     updated_eigvals_list: TensorList = []
-    for ind, (kronecker_factor, eigenbasis) in enumerate(zip(kronecker_factor_list, eigenbasis_list, strict=True)):
+    for kronecker_factor, eigenbasis in zip(kronecker_factor_list, eigenbasis_list, strict=True):
         Q = eig_utils.orthogonal_iteration(
             kronecker_factor=kronecker_factor,
             eigenbasis=eigenbasis,
             power_iter_steps=power_iter_steps,
         )
         with utils.fp32_matmul_precision("highest"):
-            approx_eigvals = eig_utils.conjugate(kronecker_factor, Q, diag=True)
-        approx_eigvals, sort_idx = torch.sort(approx_eigvals, descending=True)
-        updated_eigvals_list.append(approx_eigvals)
-        updated_eigenbasis_list.append(Q[:, sort_idx])
-        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
+            updated_eigvals_list.append(eig_utils.conjugate(kronecker_factor, Q, diag=True))
+        updated_eigenbasis_list.append(Q)
 
-    return updated_eigvals_list, updated_eigenbasis_list, exp_avg_sq
+    return updated_eigvals_list, updated_eigenbasis_list
diff --git a/emerging_optimizers/utils/eig.py b/emerging_optimizers/utils/eig.py
@@ -76,10 +76,9 @@ def orthogonal_iteration(
     """Refines an eigenbasis via power iteration with QR re-orthogonalization.
 
     Performs ``power_iter_steps`` rounds of ``Q = QR(kronecker_factor @ Q)`` starting from
-    ``eigenbasis``. The columns of ``eigenbasis`` are expected to already be approximately aligned
-    with the descending-eigenvalue ordering of ``kronecker_factor`` (maintained in the SOAP pipeline
-    by :func:`emerging_optimizers.soap.soap_utils.get_eigenbasis_qr`, which returns its eigenbases in
-    descending order).
+    ``eigenbasis``. The columns of ``eigenbasis`` are expected to already be aligned with the
+    intended descending-eigenvalue ordering of ``kronecker_factor`` (see
+    :func:`emerging_optimizers.soap.soap_utils.permute_eigenbasis_and_exp_avg_sq`).
 
     Args:
         kronecker_factor: Kronecker factor matrix (symmetric, used as the projector).
diff --git a/tests/test_soap.py b/tests/test_soap.py
@@ -582,13 +582,26 @@ def test_eigenbasis_matches_reference(self, shape: tuple, num_steps: int):
                 rtol=1e-5,
             )
 
-            for eigenbasis_test, eigenbasis_ref in zip([test_state["Q_L"], test_state["Q_R"]], ref_state["Q"]):
+            # SOAP returns eigenbases sorted by descending approximate eigenvalues while the reference
+            # keeps its power-iteration column order, so sort both by their approximate eigenvalues
+            # (against the shared kronecker factor, asserted equal above) before comparing.
+            for kronecker_factor, eigenbasis_test, eigenbasis_ref in zip(
+                [test_state["L"], test_state["R"]],
+                [test_state["Q_L"], test_state["Q_R"]],
+                ref_state["Q"],
+                strict=True,
+            ):
+                sorted_eigenbasis_list = []
+                for eigenbasis in (eigenbasis_test, eigenbasis_ref):
+                    approx_eigvals = torch.diag(eigenbasis.T @ kronecker_factor @ eigenbasis)
+                    sort_idx = torch.argsort(approx_eigvals, descending=True)
+                    sorted_eigenbasis_list.append(eigenbasis[:, sort_idx])
                 torch.testing.assert_close(
-                    eigenbasis_test,
-                    eigenbasis_ref,
+                    sorted_eigenbasis_list[0],
+                    sorted_eigenbasis_list[1],
                     atol=1e-4,
                     rtol=1e-4,
-                    msg=lambda msg: f"Eigenbasis mismatch at step {step}:\n{msg}",
+                    msg=lambda msg, step=step: f"Eigenbasis mismatch at step {step}:\n{msg}",
                 )
 
             # Compare step counters
diff --git a/tests/test_soap_utils.py b/tests/test_soap_utils.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import torch
-from _comparison import assert_close_to_identity
+from _comparison import assert_close_to_identity, assert_equal
 from absl import flags, logging
 from absl.testing import absltest, parameterized
 
@@ -61,12 +61,9 @@ def test_get_eigenbasis_qr(self, N: int, M: int) -> None:
             torch.randn(M, M, device=self.device),
             torch.randn(N, N, device=self.device),
         ]
-        exp_avg_sq = torch.abs(torch.randn(M, N, device=self.device))
-
-        eigvals_list, Q_new_list, permuted_exp_avg_sq = soap_utils.get_eigenbasis_qr(
+        eigvals_list, Q_new_list = soap_utils.get_eigenbasis_qr(
             kronecker_factor_list=kronecker_factor_list,
             eigenbasis_list=eigenbasis_list,
-            exp_avg_sq=exp_avg_sq,
             power_iter_steps=1,
         )
 
@@ -85,8 +82,6 @@ def test_get_eigenbasis_qr(self, N: int, M: int) -> None:
         self.assertEqual(len(eigvals_list), 2)
         for eigvals, kronecker_factor, Q_new in zip(eigvals_list, kronecker_factor_list, Q_new_list, strict=True):
             self.assertEqual(eigvals.shape, (kronecker_factor.shape[0],))
-            # Eigenvalues (and matching eigenbasis columns) are returned in descending order
-            self.assertTrue(torch.all(eigvals[:-1] >= eigvals[1:]))
             torch.testing.assert_close(
                 eigvals,
                 torch.diag(Q_new.t() @ kronecker_factor @ Q_new),
@@ -95,7 +90,54 @@ def test_get_eigenbasis_qr(self, N: int, M: int) -> None:
                 msg=lambda msg: f"eigvals do not match the Rayleigh quotients diag(Q^T K Q)\n\n{msg}",
             )
 
-        self.assertEqual(permuted_exp_avg_sq.shape, (M, N))
+    @parameterized.parameters(  # type: ignore[misc]
+        {"N": 4, "M": 8},
+        {"N": 16, "M": 8},
+    )
+    def test_permute_eigenbasis_and_exp_avg_sq(self, N: int, M: int) -> None:
+        """Permutes eigenbasis columns and exp_avg_sq slots consistently, in descending eigval order."""
+        g = torch.randint(-5, 6, (M, N), device=self.device) / 16.0
+        # Add a small ridge so K is full-rank (g @ g.t() is rank-deficient when M > N, etc.),
+        # which keeps the approximate eigenvalues away from numerical zero where ordering becomes
+        # ambiguous under float rounding.
+        eps_eye = lambda n: 1e-3 * torch.eye(n, device=self.device)
+        kronecker_factor_list = [g @ g.t() + eps_eye(M), g.t() @ g + eps_eye(N)]
+        eigenbasis_list = [torch.linalg.qr(Q).Q for Q in kronecker_factor_list]
+        exp_avg_sq = torch.abs(torch.randint(-5, 6, (M, N), device=self.device) / 16.0)
+
+        permuted_eigenbasis_list, permuted_exp_avg_sq = soap_utils.permute_eigenbasis_and_exp_avg_sq(
+            kronecker_factor_list,
+            eigenbasis_list,
+            exp_avg_sq,
+        )
+
+        # Compute the expected per-axis permutations from the originals.
+        sort_idx_list = []
+        for K, Q in zip(kronecker_factor_list, eigenbasis_list, strict=True):
+            sort_idx_list.append(torch.argsort(torch.diag(Q.t() @ K @ Q), descending=True))
+
+        # Each eigenbasis is column-permuted by its own sort_idx.
+        for i, (Q_old, Q_permuted) in enumerate(zip(eigenbasis_list, permuted_eigenbasis_list, strict=True)):
+            assert_equal(
+                Q_permuted,
+                Q_old[:, sort_idx_list[i]],
+                msg=lambda m, i=i: f"eigenbasis i={i} not permuted by sort_idx\n\n{m}",
+            )
+
+        # exp_avg_sq is permuted along every axis cumulatively.
+        expected_sq = exp_avg_sq
+        for i, sort_idx in enumerate(sort_idx_list):
+            expected_sq = expected_sq.index_select(i, sort_idx)
+        assert_equal(
+            permuted_exp_avg_sq,
+            expected_sq,
+            msg=lambda m: f"exp_avg_sq not permuted to match permuted eigenbases\n\n{m}",
+        )
+
+        # Permuted eigenbases yield descending approximate eigenvalues.
+        for K, Q in zip(kronecker_factor_list, permuted_eigenbasis_list, strict=True):
+            permuted_eigvals = torch.diag(Q.t() @ K @ Q)
+            self.assertTrue(torch.all(permuted_eigvals[:-1] >= permuted_eigvals[1:]))
 
     @parameterized.parameters(  # type: ignore[misc]
         {"dims": [128, 512]},
