From 0941b95d9d498bba2f6c90a8f188fcc0c09bb5f2 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 13:36:50 -0700 Subject: [PATCH 01/10] feat: add scaled CANS inverse root utility Signed-off-by: mkhona --- .../soap/matrix_root_inverse_utils.py | 101 ++++++++++++++++++ tests/test_matrix_root_inverse_utils.py | 77 +++++++++++++ 2 files changed, 178 insertions(+) create mode 100644 emerging_optimizers/soap/matrix_root_inverse_utils.py create mode 100644 tests/test_matrix_root_inverse_utils.py diff --git a/emerging_optimizers/soap/matrix_root_inverse_utils.py b/emerging_optimizers/soap/matrix_root_inverse_utils.py new file mode 100644 index 0000000..1e18cad --- /dev/null +++ b/emerging_optimizers/soap/matrix_root_inverse_utils.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch +from torch import Tensor + +from emerging_optimizers import utils +from emerging_optimizers.utils import FP32MatmulPrecT + + +__all__ = ["scaled_cans_coupled_ns"] + +# All constant arithmetic, including the 1.01 safety factor, is folded into these coefficients. +_CANS_COEFFS = ( + (5.182503604966906, -5.126830178299687), + (2.586120737395915, -0.641538812403133), + (2.567364126726186, -0.6391058222170474), + (2.520560084348265, -0.6330225823828756), + (2.410759275435182, -0.6186815444268036), + (2.1883348130094173, -0.5893091162177136), + (1.8595760874873613, -0.5449991062102938), + (1.589020160467417, -0.5075811685214573), + (1.5051653981684994, -0.4957799077972079), + (1.4925557853149838, -0.49259266842078675), +) + + +def scaled_cans_coupled_ns( + x: Tensor, + eps: float = 1e-12, + fp32_matmul_prec: FP32MatmulPrecT = "highest", +) -> Tensor: + """Compute inverse square roots with scaled coupled CANS Newton-Schulz. + + CANS polynomial-based inverse-root computation from https://arxiv.org/abs/2506.10935. + + This implementation applies the CANS orthogonalization polynomials to the coupled + Newton-Schulz iteration for a symmetric positive-definite matrix. It uses a fixed + ten-step schedule and normalizes with the matrix infinity norm rather than the exact + spectral norm. The infinity norm is an inexpensive upper bound, but can conservatively + scale matrices whose rows contain substantial cancellation and consequently slow + convergence for their smallest eigenvalues. + + The tabulated coefficients fold a 1% spectral safety margin into the polynomial. Starting + from the unscaled CANS pairs ``(beta, alpha)``, steps zero through eight keep ``beta`` and + use ``alpha / 1.01``. The final pair additionally absorbs the output normalization and is + computed as ``(beta / sqrt(1.01), alpha / 1.01**1.5)``. This is algebraically equivalent + in exact arithmetic to normalizing by ``1.01 * inf_norm`` and applying the original + coefficients. The literals were generated with Python IEEE-754 binary64 arithmetic and + rounded to the nearest representable binary64 value. In particular, + ``1.5 / sqrt(1.01)`` evaluates to ``1.4925557853149838``; the decimal + ``1.492555785314984`` is one binary64 ULP higher. All constant arithmetic is folded into + these literals. At runtime they are applied to FP32 tensors, except that ``"medium"`` + explicitly casts the iteration to BF16, matching Muon's Newton-Schulz implementation. + + In practice, the result is approximate because the iteration is truncated after ten steps, + the normalization may overestimate the largest eigenvalue, matrix multiplications may use + reduced precision, and the returned inverse root is explicitly symmetrized. ``eps`` also + clamps the normalization for degenerate inputs. + + Args: + x: A 2D symmetric positive-definite FP32 matrix or 3D batch of matrices. + eps: Lower bound used when normalizing the matrices. + fp32_matmul_prec: Precision used for FP32 matrix multiplications: ``"medium"`` for BF16, + ``"high"`` for TF32, or ``"highest"`` for FP32. + + Returns: + The approximate inverse square root as an FP32 tensor with the same shape as ``x``. + """ + if x.dim() not in (2, 3) or x.shape[-2] != x.shape[-1]: + raise TypeError(f"x must be a square matrix or batch of square matrices, got shape {tuple(x.shape)}") + if x.dtype != torch.float32: + raise TypeError(f"x must be in float32, got {x.dtype}") + + with utils.fp32_matmul_precision(fp32_matmul_prec): + inf_norm = torch.linalg.matrix_norm(x, ord=float("inf"), dim=(-2, -1), keepdim=True).clamp_min_(eps) + y = x / inf_norm + if torch.get_float32_matmul_precision() == "medium": + y = y.to(torch.bfloat16) + + z = torch.eye(x.shape[-1], device=x.device, dtype=y.dtype).expand_as(y) + cans_addmm = torch.addmm if x.dim() == 2 else torch.baddbmm + for beta, alpha in _CANS_COEFFS: + p = z @ y + z = cans_addmm(z, p, z, beta=beta, alpha=alpha) + y = cans_addmm(y, y, p, beta=beta, alpha=alpha) + + z = z.to(torch.float32) + z.mul_(torch.rsqrt(inf_norm)) + return (z + z.mT) / 2.0 diff --git a/tests/test_matrix_root_inverse_utils.py b/tests/test_matrix_root_inverse_utils.py new file mode 100644 index 0000000..a01f141 --- /dev/null +++ b/tests/test_matrix_root_inverse_utils.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch +from absl import flags, logging +from absl.testing import absltest, parameterized + +from emerging_optimizers.soap.matrix_root_inverse_utils import scaled_cans_coupled_ns +from emerging_optimizers.utils import FP32MatmulPrecT + + +flags.DEFINE_enum("device", "cpu", ["cpu", "cuda"], "Device to run tests on") +flags.DEFINE_integer("seed", None, "Random seed for reproducible tests") +FLAGS = flags.FLAGS + + +def setUpModule() -> None: + if FLAGS.seed is not None: + logging.info("Setting random seed to %d", FLAGS.seed) + torch.manual_seed(FLAGS.seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(FLAGS.seed) + + +class MatrixRootInverseUtilsTest(parameterized.TestCase): + @parameterized.product( + shape=[(4, 4), (2, 4, 4)], + fp32_matmul_prec=["medium", "high", "highest"], + ) + def test_scaled_cans_smoke( + self, + shape: tuple[int, ...], + fp32_matmul_prec: FP32MatmulPrecT, + ) -> None: + x = torch.randn(*shape, device=FLAGS.device) + matrix = x @ x.mT + 0.1 * torch.eye(shape[-1], device=FLAGS.device) + previous_precision = torch.get_float32_matmul_precision() + + inverse_root = scaled_cans_coupled_ns(matrix, fp32_matmul_prec=fp32_matmul_prec) + + self.assertEqual(inverse_root.shape, matrix.shape) + self.assertEqual(inverse_root.dtype, torch.float32) + self.assertEqual(torch.get_float32_matmul_precision(), previous_precision) + + def test_scaled_cans_inverse_root_accuracy(self) -> None: + matrix = torch.tensor( + [[2.0, 0.5], [0.5, 1.5]], + device=FLAGS.device, + ) + + inverse_root = scaled_cans_coupled_ns(matrix) + identity = torch.eye(matrix.shape[-1], device=FLAGS.device) + whitened_matrix = inverse_root @ matrix @ inverse_root + matrix_root = torch.linalg.inv(inverse_root) + reconstructed_matrix = matrix_root @ matrix_root + + torch.testing.assert_close(whitened_matrix, identity, atol=2e-4, rtol=2e-4) + torch.testing.assert_close(reconstructed_matrix, matrix, atol=2e-4, rtol=2e-4) + + def test_scaled_cans_rejects_non_fp32_tensor(self) -> None: + with self.assertRaisesRegex(TypeError, "must be in float32"): + scaled_cans_coupled_ns(torch.eye(4, device=FLAGS.device, dtype=torch.bfloat16)) + + +if __name__ == "__main__": + absltest.main() From 1a36f62a8b78f2e3a19aea42b6ad6c38da7777b9 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 16:08:08 -0700 Subject: [PATCH 02/10] chore: update CANS copyright year Signed-off-by: mkhona --- emerging_optimizers/soap/matrix_root_inverse_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/emerging_optimizers/soap/matrix_root_inverse_utils.py b/emerging_optimizers/soap/matrix_root_inverse_utils.py index 1e18cad..f40c1d1 100644 --- a/emerging_optimizers/soap/matrix_root_inverse_utils.py +++ b/emerging_optimizers/soap/matrix_root_inverse_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); From 4aa6f49548fb52482d72737aad44af35840ac935 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 16:18:42 -0700 Subject: [PATCH 03/10] style: use explicit symmetrization scale Signed-off-by: mkhona --- emerging_optimizers/soap/matrix_root_inverse_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/emerging_optimizers/soap/matrix_root_inverse_utils.py b/emerging_optimizers/soap/matrix_root_inverse_utils.py index f40c1d1..35769d4 100644 --- a/emerging_optimizers/soap/matrix_root_inverse_utils.py +++ b/emerging_optimizers/soap/matrix_root_inverse_utils.py @@ -98,4 +98,4 @@ def scaled_cans_coupled_ns( z = z.to(torch.float32) z.mul_(torch.rsqrt(inf_norm)) - return (z + z.mT) / 2.0 + return (z + z.mT) * 0.5 From 603af790c5e870641e8ac542e467883203a3d040 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 16:26:20 -0700 Subject: [PATCH 04/10] refactor: broadcast batched CANS identity Signed-off-by: mkhona --- emerging_optimizers/soap/matrix_root_inverse_utils.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/emerging_optimizers/soap/matrix_root_inverse_utils.py b/emerging_optimizers/soap/matrix_root_inverse_utils.py index 35769d4..38fda43 100644 --- a/emerging_optimizers/soap/matrix_root_inverse_utils.py +++ b/emerging_optimizers/soap/matrix_root_inverse_utils.py @@ -89,12 +89,14 @@ def scaled_cans_coupled_ns( if torch.get_float32_matmul_precision() == "medium": y = y.to(torch.bfloat16) - z = torch.eye(x.shape[-1], device=x.device, dtype=y.dtype).expand_as(y) - cans_addmm = torch.addmm if x.dim() == 2 else torch.baddbmm + z = torch.eye(x.shape[-1], device=x.device, dtype=y.dtype) + if x.dim() == 3: + z = z.unsqueeze(0) + for beta, alpha in _CANS_COEFFS: p = z @ y - z = cans_addmm(z, p, z, beta=beta, alpha=alpha) - y = cans_addmm(y, y, p, beta=beta, alpha=alpha) + z = torch.add(z * beta, p @ z, alpha=alpha) + y = torch.add(y * beta, y @ p, alpha=alpha) z = z.to(torch.float32) z.mul_(torch.rsqrt(inf_norm)) From c25cb3ad9864ad082cf4187df0d7ed5ba1230eeb Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 16:33:00 -0700 Subject: [PATCH 05/10] test: cover batched CANS accuracy Signed-off-by: mkhona --- tests/test_matrix_root_inverse_utils.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/test_matrix_root_inverse_utils.py b/tests/test_matrix_root_inverse_utils.py index a01f141..429ff8d 100644 --- a/tests/test_matrix_root_inverse_utils.py +++ b/tests/test_matrix_root_inverse_utils.py @@ -53,14 +53,20 @@ def test_scaled_cans_smoke( self.assertEqual(inverse_root.dtype, torch.float32) self.assertEqual(torch.get_float32_matmul_precision(), previous_precision) - def test_scaled_cans_inverse_root_accuracy(self) -> None: - matrix = torch.tensor( + @parameterized.parameters((2, 2), (2, 2, 2)) # type: ignore[misc] + def test_scaled_cans_inverse_root_accuracy(self, shape: tuple[int, ...]) -> None: + base_matrix = torch.tensor( [[2.0, 0.5], [0.5, 1.5]], device=FLAGS.device, ) + if len(shape) == 2: + matrix = base_matrix + else: + batch_scale = torch.arange(1, shape[0] + 1, device=FLAGS.device).view(-1, 1, 1) + matrix = base_matrix.unsqueeze(0) * batch_scale inverse_root = scaled_cans_coupled_ns(matrix) - identity = torch.eye(matrix.shape[-1], device=FLAGS.device) + identity = torch.eye(matrix.shape[-1], device=FLAGS.device).expand_as(matrix) whitened_matrix = inverse_root @ matrix @ inverse_root matrix_root = torch.linalg.inv(inverse_root) reconstructed_matrix = matrix_root @ matrix_root From 64445c9399fc75ab03dba6c77672cc06d8ac9365 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 16:34:46 -0700 Subject: [PATCH 06/10] test: expand CANS accuracy coverage Signed-off-by: mkhona --- tests/test_matrix_root_inverse_utils.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/tests/test_matrix_root_inverse_utils.py b/tests/test_matrix_root_inverse_utils.py index 429ff8d..21faf05 100644 --- a/tests/test_matrix_root_inverse_utils.py +++ b/tests/test_matrix_root_inverse_utils.py @@ -13,6 +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 absl import flags, logging from absl.testing import absltest, parameterized @@ -53,25 +54,30 @@ def test_scaled_cans_smoke( self.assertEqual(inverse_root.dtype, torch.float32) self.assertEqual(torch.get_float32_matmul_precision(), previous_precision) - @parameterized.parameters((2, 2), (2, 2, 2)) # type: ignore[misc] + @parameterized.parameters((8, 8), (16, 16), (2, 8, 8), (3, 16, 16)) # type: ignore[misc] def test_scaled_cans_inverse_root_accuracy(self, shape: tuple[int, ...]) -> None: - base_matrix = torch.tensor( - [[2.0, 0.5], [0.5, 1.5]], - device=FLAGS.device, - ) + matrix_size = shape[-1] + base_matrix = 2.0 * torch.eye(matrix_size, device=FLAGS.device) + base_matrix.diagonal(offset=1).fill_(0.25) + base_matrix.diagonal(offset=-1).fill_(0.25) if len(shape) == 2: matrix = base_matrix else: - batch_scale = torch.arange(1, shape[0] + 1, device=FLAGS.device).view(-1, 1, 1) + batch_scale = torch.arange( + 1, + shape[0] + 1, + device=FLAGS.device, + dtype=base_matrix.dtype, + ).view(-1, 1, 1) matrix = base_matrix.unsqueeze(0) * batch_scale inverse_root = scaled_cans_coupled_ns(matrix) - identity = torch.eye(matrix.shape[-1], device=FLAGS.device).expand_as(matrix) whitened_matrix = inverse_root @ matrix @ inverse_root matrix_root = torch.linalg.inv(inverse_root) reconstructed_matrix = matrix_root @ matrix_root - torch.testing.assert_close(whitened_matrix, identity, atol=2e-4, rtol=2e-4) + for whitened_matrix_slice in whitened_matrix.reshape(-1, matrix_size, matrix_size): + assert_close_to_identity(whitened_matrix_slice, off_diag_atol=2e-4, diag_atol=2e-4) torch.testing.assert_close(reconstructed_matrix, matrix, atol=2e-4, rtol=2e-4) def test_scaled_cans_rejects_non_fp32_tensor(self) -> None: From 0b63790a8759d2b4aae3afccefe326e93f9c2658 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 16:36:56 -0700 Subject: [PATCH 07/10] refactor: clarify CANS inverse-root function name Signed-off-by: mkhona --- .../soap/matrix_root_inverse_utils.py | 4 ++-- tests/test_matrix_root_inverse_utils.py | 14 +++++++------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/emerging_optimizers/soap/matrix_root_inverse_utils.py b/emerging_optimizers/soap/matrix_root_inverse_utils.py index 38fda43..68cf8fa 100644 --- a/emerging_optimizers/soap/matrix_root_inverse_utils.py +++ b/emerging_optimizers/soap/matrix_root_inverse_utils.py @@ -19,7 +19,7 @@ from emerging_optimizers.utils import FP32MatmulPrecT -__all__ = ["scaled_cans_coupled_ns"] +__all__ = ["mat_root_inv_via_scaled_cans"] # All constant arithmetic, including the 1.01 safety factor, is folded into these coefficients. _CANS_COEFFS = ( @@ -36,7 +36,7 @@ ) -def scaled_cans_coupled_ns( +def mat_root_inv_via_scaled_cans( x: Tensor, eps: float = 1e-12, fp32_matmul_prec: FP32MatmulPrecT = "highest", diff --git a/tests/test_matrix_root_inverse_utils.py b/tests/test_matrix_root_inverse_utils.py index 21faf05..e96e971 100644 --- a/tests/test_matrix_root_inverse_utils.py +++ b/tests/test_matrix_root_inverse_utils.py @@ -17,7 +17,7 @@ from absl import flags, logging from absl.testing import absltest, parameterized -from emerging_optimizers.soap.matrix_root_inverse_utils import scaled_cans_coupled_ns +from emerging_optimizers.soap.matrix_root_inverse_utils import mat_root_inv_via_scaled_cans from emerging_optimizers.utils import FP32MatmulPrecT @@ -39,7 +39,7 @@ class MatrixRootInverseUtilsTest(parameterized.TestCase): shape=[(4, 4), (2, 4, 4)], fp32_matmul_prec=["medium", "high", "highest"], ) - def test_scaled_cans_smoke( + def test_mat_root_inv_via_scaled_cans_smoke( self, shape: tuple[int, ...], fp32_matmul_prec: FP32MatmulPrecT, @@ -48,14 +48,14 @@ def test_scaled_cans_smoke( matrix = x @ x.mT + 0.1 * torch.eye(shape[-1], device=FLAGS.device) previous_precision = torch.get_float32_matmul_precision() - inverse_root = scaled_cans_coupled_ns(matrix, fp32_matmul_prec=fp32_matmul_prec) + inverse_root = mat_root_inv_via_scaled_cans(matrix, fp32_matmul_prec=fp32_matmul_prec) self.assertEqual(inverse_root.shape, matrix.shape) self.assertEqual(inverse_root.dtype, torch.float32) self.assertEqual(torch.get_float32_matmul_precision(), previous_precision) @parameterized.parameters((8, 8), (16, 16), (2, 8, 8), (3, 16, 16)) # type: ignore[misc] - def test_scaled_cans_inverse_root_accuracy(self, shape: tuple[int, ...]) -> None: + def test_mat_root_inv_via_scaled_cans_accuracy(self, shape: tuple[int, ...]) -> None: matrix_size = shape[-1] base_matrix = 2.0 * torch.eye(matrix_size, device=FLAGS.device) base_matrix.diagonal(offset=1).fill_(0.25) @@ -71,7 +71,7 @@ def test_scaled_cans_inverse_root_accuracy(self, shape: tuple[int, ...]) -> None ).view(-1, 1, 1) matrix = base_matrix.unsqueeze(0) * batch_scale - inverse_root = scaled_cans_coupled_ns(matrix) + inverse_root = mat_root_inv_via_scaled_cans(matrix) whitened_matrix = inverse_root @ matrix @ inverse_root matrix_root = torch.linalg.inv(inverse_root) reconstructed_matrix = matrix_root @ matrix_root @@ -80,9 +80,9 @@ def test_scaled_cans_inverse_root_accuracy(self, shape: tuple[int, ...]) -> None assert_close_to_identity(whitened_matrix_slice, off_diag_atol=2e-4, diag_atol=2e-4) torch.testing.assert_close(reconstructed_matrix, matrix, atol=2e-4, rtol=2e-4) - def test_scaled_cans_rejects_non_fp32_tensor(self) -> None: + def test_mat_root_inv_via_scaled_cans_rejects_non_fp32_tensor(self) -> None: with self.assertRaisesRegex(TypeError, "must be in float32"): - scaled_cans_coupled_ns(torch.eye(4, device=FLAGS.device, dtype=torch.bfloat16)) + mat_root_inv_via_scaled_cans(torch.eye(4, device=FLAGS.device, dtype=torch.bfloat16)) if __name__ == "__main__": From 673fb97c8e79df702eda4bcf1d3dce964792003f Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 13:37:06 -0700 Subject: [PATCH 08/10] feat: add Online KL-Shampoo optimizer Signed-off-by: mkhona --- docs/apidocs/soap.md | 18 ++ emerging_optimizers/soap/__init__.py | 2 + emerging_optimizers/soap/okls.py | 268 +++++++++++++++++++++++++++ tests/test_okls.py | 73 ++++++++ tests/test_registry.py | 1 + 5 files changed, 362 insertions(+) create mode 100644 emerging_optimizers/soap/okls.py create mode 100644 tests/test_okls.py diff --git a/docs/apidocs/soap.md b/docs/apidocs/soap.md index b123568..406b499 100644 --- a/docs/apidocs/soap.md +++ b/docs/apidocs/soap.md @@ -25,6 +25,17 @@ emerging_optimizers.soap .. autofunction:: update_eigenbasis_and_exp_avgs +:hidden:`OKLS` +~~~~~~~~~~~~~~~ + +.. currentmodule:: emerging_optimizers.soap.okls + +.. autoclass:: OKLS + :members: + +.. autofunction:: update_kronecker_factors_okls + + :hidden:`REKLS` ~~~~~~~~~~~~~~~ @@ -39,4 +50,11 @@ emerging_optimizers.soap.soap_utils .. automodule:: emerging_optimizers.soap.soap_utils :members: + + +emerging_optimizers.soap.matrix_root_inverse_utils +===================================== + +.. automodule:: emerging_optimizers.soap.matrix_root_inverse_utils + :members: ``` diff --git a/emerging_optimizers/soap/__init__.py b/emerging_optimizers/soap/__init__.py index 7cbe99b..859b7a3 100644 --- a/emerging_optimizers/soap/__init__.py +++ b/emerging_optimizers/soap/__init__.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from emerging_optimizers.soap.moso import MOSO +from emerging_optimizers.soap.okls import OKLS from emerging_optimizers.soap.rekls import REKLS from emerging_optimizers.soap.soap import SOAP from emerging_optimizers.soap.stacked_soap import StackedSoap @@ -20,6 +21,7 @@ __all__ = [ "MOSO", + "OKLS", "REKLS", "SOAP", "StackedSoap", diff --git a/emerging_optimizers/soap/okls.py b/emerging_optimizers/soap/okls.py new file mode 100644 index 0000000..bc342cd --- /dev/null +++ b/emerging_optimizers/soap/okls.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math +from typing import TYPE_CHECKING, Callable, override + + +if TYPE_CHECKING: + from typing import overload + +import torch +from torch import optim +from torch.optim.optimizer import ParamsT + +from emerging_optimizers import mixin as opt_mixin +from emerging_optimizers import registry +from emerging_optimizers.soap.matrix_root_inverse_utils import scaled_cans_coupled_ns +from emerging_optimizers.utils import FP32MatmulPrecT + + +__all__ = ["OKLS", "update_kronecker_factors_okls"] + + +def _update_inverse_roots( + kronecker_factor_list: list[torch.Tensor], + inverse_root_list: list[torch.Tensor], + ridge_eps: float, + cans_fp32_matmul_prec: FP32MatmulPrecT, +) -> None: + for kronecker_factor, inverse_root in zip(kronecker_factor_list, inverse_root_list, strict=True): + inverse_root.copy_( + scaled_cans_coupled_ns( + kronecker_factor, + eps=ridge_eps, + fp32_matmul_prec=cans_fp32_matmul_prec, + ) + ) + + +def _initialize_preconditioners( + kronecker_factor_list: list[torch.Tensor], + inverse_root_list: list[torch.Tensor], + grad: torch.Tensor, + ridge_eps: float, + cans_fp32_matmul_prec: FP32MatmulPrecT, +) -> None: + rows, cols = grad.shape + grad_norm_sq = grad.square().sum() + factor_left, factor_right = kronecker_factor_list + + factor_left.copy_(grad @ grad.T) + factor_left.mul_(torch.sqrt(rows / (cols * grad_norm_sq + ridge_eps))) + factor_left.copy_((factor_left + factor_left.T) / 2.0) + diagonal_shift_left = torch.linalg.norm(factor_left) / math.sqrt(rows) + factor_left.diagonal().add_(diagonal_shift_left + ridge_eps) + + factor_right.copy_(grad.T @ grad) + factor_right.mul_(torch.sqrt(cols / (rows * grad_norm_sq + ridge_eps))) + factor_right.copy_((factor_right + factor_right.T) / 2.0) + diagonal_shift_right = torch.linalg.norm(factor_right) / math.sqrt(cols) + factor_right.diagonal().add_(diagonal_shift_right + ridge_eps) + + _update_inverse_roots( + kronecker_factor_list, + inverse_root_list, + ridge_eps, + cans_fp32_matmul_prec, + ) + + +@torch.no_grad() # type: ignore[misc] +def update_kronecker_factors_okls( + kronecker_factor_list: list[torch.Tensor], + inverse_root_list: list[torch.Tensor], + grad: torch.Tensor, + shampoo_beta: float, + ridge_eps: float, +) -> None: + """Update KL-Shampoo factors using the previous inverse-square-root preconditioners. + + Args: + kronecker_factor_list: Left and right covariance factors. + inverse_root_list: Previous inverse square roots of the left and right factors. + grad: Matrix gradient. + shampoo_beta: EMA coefficient for the factors. + ridge_eps: Diagonal stability offset. + """ + if grad.dim() != 2: + raise TypeError("OKLS is only supported for 2D tensors") + + factor_left, factor_right = kronecker_factor_list + inverse_root_left, inverse_root_right = inverse_root_list + rows, cols = grad.shape + + grad_right_preconditioned = grad @ inverse_root_right + factor_left.lerp_(grad_right_preconditioned @ grad_right_preconditioned.T / cols, 1 - shampoo_beta) + factor_left.copy_((factor_left + factor_left.T) / 2.0) + factor_left.diagonal().add_(ridge_eps) + + grad_left_preconditioned = inverse_root_left @ grad + factor_right.lerp_(grad_left_preconditioned.T @ grad_left_preconditioned / rows, 1 - shampoo_beta) + factor_right.copy_((factor_right + factor_right.T) / 2.0) + factor_right.diagonal().add_(ridge_eps) + + +@registry.register_optimizer("okls") +class OKLS(opt_mixin.WeightDecayMixin, optim.Optimizer): + """Online KL-Shampoo with scaled CANS inverse roots and zero-staleness preconditioning. + + Args: + params: Iterable of 2D CUDA parameters to optimize or dicts defining parameter groups. + lr: Learning rate. + beta1: Nesterov momentum EMA coefficient. + beta2: KL-Shampoo factor EMA coefficient. + ridge_eps: Numerical stability offset added to the KL-Shampoo factors. + weight_decay: PaLM weight-decay coefficient. + cans_fp32_matmul_prec: Precision used for FP32 matrix multiplications in CANS: ``"medium"`` for BF16, + ``"high"`` for TF32, or ``"highest"`` for FP32. + """ + + def __init__( + self, + params: ParamsT, + *, + lr: float, + beta1: float = 0.9684, + beta2: float = 0.9482, + ridge_eps: float = 1e-9, + weight_decay: float = 0.0, + cans_fp32_matmul_prec: FP32MatmulPrecT = "high", + ) -> None: + self.weight_decay_method = "palm" + self.cans_fp32_matmul_prec = cans_fp32_matmul_prec + + if lr < 0.0: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= beta1 < 1.0: + raise ValueError(f"Invalid beta1: {beta1}") + if not 0.0 <= beta2 < 1.0: + raise ValueError(f"Invalid beta2: {beta2}") + if ridge_eps < 0.0: + raise ValueError(f"Invalid ridge epsilon: {ridge_eps}") + if weight_decay < 0.0: + raise ValueError(f"Invalid weight_decay: {weight_decay}") + + defaults = { + "lr": lr, + "beta1": beta1, + "beta2": beta2, + "ridge_eps": ridge_eps, + "weight_decay": weight_decay, + } + super().__init__(params, defaults) + + @torch.no_grad() # type: ignore[misc] + def _init_group( + self, + group: dict, + skip_non_grad_params: bool = True, + ) -> None: + for p in group["params"]: + if skip_non_grad_params and p.grad is None: + continue + + if p.dim() != 2: + raise TypeError("OKLS is only supported for 2D tensors") + if not p.is_cuda: + raise TypeError("OKLS only supports CUDA tensors") + + state = self.state[p] + if len(state) == 0: + state["step"] = 0 + state["exp_avg"] = torch.zeros_like(p, dtype=torch.float32) + state["L"] = p.new_zeros((p.shape[0], p.shape[0]), dtype=torch.float32) + state["R"] = p.new_zeros((p.shape[1], p.shape[1]), dtype=torch.float32) + state["P_L"] = p.new_zeros((p.shape[0], p.shape[0]), dtype=torch.float32) + state["P_R"] = p.new_zeros((p.shape[1], p.shape[1]), dtype=torch.float32) + + if TYPE_CHECKING: + + @overload + def step(self, closure: None = ...) -> None: ... + + @overload + def step(self, closure: Callable[[], float]) -> float: ... + + @torch.no_grad() # type: ignore[misc] + @override + def step(self, closure: Callable[[], float] | None = None) -> float | None: + """Perform a single optimization step. + + Args: + closure: Unsupported; must be ``None``. + """ + if closure is not None: + raise ValueError("closure is not supported") + + for group in self.param_groups: + self._init_group(group) + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue # pragma: no cover + + grad = p.grad.to(torch.float32) + state = self.state[p] + kronecker_factor_list = [state["L"], state["R"]] + inverse_root_list = [state["P_L"], state["P_R"]] + ridge_eps = group["ridge_eps"] + + if state["step"] == 0: + _initialize_preconditioners( + kronecker_factor_list, + inverse_root_list, + grad, + ridge_eps, + self.cans_fp32_matmul_prec, + ) + + beta1 = group["beta1"] + state["exp_avg"].lerp_(grad, 1 - beta1) + nesterov_momentum = torch.lerp(grad, state["exp_avg"], beta1) + + update_kronecker_factors_okls( + kronecker_factor_list=kronecker_factor_list, + inverse_root_list=inverse_root_list, + grad=grad, + shampoo_beta=group["beta2"], + ridge_eps=ridge_eps, + ) + _update_inverse_roots( + kronecker_factor_list, + inverse_root_list, + ridge_eps, + self.cans_fp32_matmul_prec, + ) + + preconditioned_update = inverse_root_list[0] @ nesterov_momentum @ inverse_root_list[1] + rows, cols = grad.shape + nesterov_variance = ((1 - beta1) / (1 + beta1)) * (1 + 2 * beta1 - 2 * beta1**3) + momentum_scale = nesterov_variance**-0.5 + shape_scale = math.sqrt(rows / cols) / (math.sqrt(rows) + math.sqrt(cols)) + + self._apply_weight_decay_inplace( + p, + grad, + group["lr"], + group["weight_decay"], + ) + p.add_( + preconditioned_update.to(p.dtype), + alpha=-group["lr"] * momentum_scale * shape_scale, + ) + state["step"] += 1 + + return None diff --git a/tests/test_okls.py b/tests/test_okls.py new file mode 100644 index 0000000..0efd0c2 --- /dev/null +++ b/tests/test_okls.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch +from absl import flags, logging +from absl.testing import absltest, parameterized + +from emerging_optimizers.soap.okls import OKLS + + +flags.DEFINE_enum("device", "cpu", ["cpu", "cuda"], "Device to run tests on") +flags.DEFINE_integer("seed", None, "Random seed for reproducible tests") +FLAGS = flags.FLAGS + + +def setUpModule() -> None: + if FLAGS.seed is not None: + logging.info("Setting random seed to %d", FLAGS.seed) + torch.manual_seed(FLAGS.seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(FLAGS.seed) + + +class OKLSTest(parameterized.TestCase): + def test_step_initializes_state_and_updates_parameter(self) -> None: + if FLAGS.device != "cuda": + self.skipTest("OKLS requires CUDA") + + param = torch.nn.Parameter(torch.randn(4, 3, device=FLAGS.device)) + original = param.detach().clone() + param.grad = torch.randn_like(param) + optimizer = OKLS([param], lr=0.01, ridge_eps=1e-9) + + optimizer.step() + + self.assertEqual(optimizer.weight_decay_method, "palm") + self.assertEqual(optimizer.cans_fp32_matmul_prec, "high") + self.assertNotIn("lr_peak", optimizer.param_groups[0]) + self.assertFalse(torch.equal(param, original)) + self.assertTrue(torch.isfinite(param).all()) + state = optimizer.state[param] + self.assertEqual(state["step"], 1) + self.assertCountEqual(state.keys(), ["step", "exp_avg", "L", "R", "P_L", "P_R"]) + self.assertEqual(state["L"].shape, (4, 4)) + self.assertEqual(state["R"].shape, (3, 3)) + self.assertEqual(state["P_L"].shape, (4, 4)) + self.assertEqual(state["P_R"].shape, (3, 3)) + for value in state.values(): + if isinstance(value, torch.Tensor): + self.assertEqual(value.dtype, torch.float32) + self.assertTrue(torch.isfinite(value).all()) + + def test_non_2d_parameter_raises_type_error(self) -> None: + param = torch.nn.Parameter(torch.randn(2, 4, 3, device=FLAGS.device)) + param.grad = torch.randn_like(param) + optimizer = OKLS([param], lr=0.01) + with self.assertRaisesRegex(TypeError, "only supported for 2D"): + optimizer.step() + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/test_registry.py b/tests/test_registry.py index 71e805c..0828cc7 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -60,6 +60,7 @@ def __init__(self, params, lr=0.01): ("psgd_pro", psgd.PSGDPro), ("scion", scion.Scion), ("soap", soap.SOAP), + ("okls", soap.OKLS), ("lion", scalar_optimizers.Lion), ("laprop", scalar_optimizers.LaProp), ) From 379ba77f0d454263b8b86529d4463d18860cef49 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 16:37:21 -0700 Subject: [PATCH 09/10] refactor: update OKLS CANS inverse-root call --- emerging_optimizers/soap/okls.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/emerging_optimizers/soap/okls.py b/emerging_optimizers/soap/okls.py index bc342cd..6fcf7df 100644 --- a/emerging_optimizers/soap/okls.py +++ b/emerging_optimizers/soap/okls.py @@ -25,7 +25,7 @@ from emerging_optimizers import mixin as opt_mixin from emerging_optimizers import registry -from emerging_optimizers.soap.matrix_root_inverse_utils import scaled_cans_coupled_ns +from emerging_optimizers.soap.matrix_root_inverse_utils import mat_root_inv_via_scaled_cans from emerging_optimizers.utils import FP32MatmulPrecT @@ -40,7 +40,7 @@ def _update_inverse_roots( ) -> None: for kronecker_factor, inverse_root in zip(kronecker_factor_list, inverse_root_list, strict=True): inverse_root.copy_( - scaled_cans_coupled_ns( + mat_root_inv_via_scaled_cans( kronecker_factor, eps=ridge_eps, fp32_matmul_prec=cans_fp32_matmul_prec, From 42cf97fb5f70bfcceeca4bbf0fcca1b1fcb18e90 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 28 Jul 2026 17:01:16 -0700 Subject: [PATCH 10/10] fix: use decoupled weight decay in OKLS --- emerging_optimizers/soap/okls.py | 4 ++-- tests/test_okls.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/emerging_optimizers/soap/okls.py b/emerging_optimizers/soap/okls.py index 6fcf7df..14ac085 100644 --- a/emerging_optimizers/soap/okls.py +++ b/emerging_optimizers/soap/okls.py @@ -124,7 +124,7 @@ class OKLS(opt_mixin.WeightDecayMixin, optim.Optimizer): beta1: Nesterov momentum EMA coefficient. beta2: KL-Shampoo factor EMA coefficient. ridge_eps: Numerical stability offset added to the KL-Shampoo factors. - weight_decay: PaLM weight-decay coefficient. + weight_decay: Decoupled weight-decay coefficient. cans_fp32_matmul_prec: Precision used for FP32 matrix multiplications in CANS: ``"medium"`` for BF16, ``"high"`` for TF32, or ``"highest"`` for FP32. """ @@ -140,7 +140,7 @@ def __init__( weight_decay: float = 0.0, cans_fp32_matmul_prec: FP32MatmulPrecT = "high", ) -> None: - self.weight_decay_method = "palm" + self.weight_decay_method = "decoupled" self.cans_fp32_matmul_prec = cans_fp32_matmul_prec if lr < 0.0: diff --git a/tests/test_okls.py b/tests/test_okls.py index 0efd0c2..e1430ee 100644 --- a/tests/test_okls.py +++ b/tests/test_okls.py @@ -44,7 +44,7 @@ def test_step_initializes_state_and_updates_parameter(self) -> None: optimizer.step() - self.assertEqual(optimizer.weight_decay_method, "palm") + self.assertEqual(optimizer.weight_decay_method, "decoupled") self.assertEqual(optimizer.cans_fp32_matmul_prec, "high") self.assertNotIn("lr_peak", optimizer.param_groups[0]) self.assertFalse(torch.equal(param, original))