OKLS optimizer

[mkhona-nvidia]

https://blog.tilderesearch.com/blog/online-kl-shampoo

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Zero epsilon produces NaN state

When ridge_eps=0 and the first gradient is all zero, initialization computes sqrt(rows / 0) and multiplies the zero covariance by infinity, causing NaN factors, inverse roots, and parameter values on the first step.

Suggested change

	if ridge_eps < 0.0:
	raise ValueError(f"Invalid ridge epsilon: {ridge_eps}")
	if ridge_eps <= 0.0:
	raise ValueError(f"Invalid ridge epsilon: {ridge_eps}")

[greptile-apps]

Factor matmuls inherit global precision

The covariance and final preconditioning matmuls execute outside fp32_matmul_precision, so a process-wide medium setting silently reduces their precision and makes persistent optimizer factors depend on unrelated global configuration.

Knowledge Base Used: SOAP: Shampoo-style Preconditioning

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Greptile Summary

Adds the OKLS optimizer and its scaled-CANS inverse-root implementation.

• Registers and publicly exports OKLS.
• Adds API documentation for OKLS and the matrix inverse-root utility.
• Adds optimizer, numerical utility, and registry tests.

Confidence Score: 3/5

The PR is not yet safe to merge because zero epsilon can corrupt optimizer state with NaNs and configured matmul precision is not applied throughout the OKLS update.

The current implementation still permits a zero epsilon that makes zero-gradient initialization evaluate zero times infinity, while factor and final preconditioning matmuls remain controlled by unrelated process-wide precision.

Files Needing Attention: emerging_optimizers/soap/okls.py

Important Files Changed

Filename	Overview
emerging_optimizers/soap/okls.py	Implements OKLS state initialization, factor updates, inverse-root refreshes, momentum, and parameter updates.
emerging_optimizers/soap/matrix_root_inverse_utils.py	Implements the fixed-schedule scaled-CANS approximation for FP32 matrix inverse square roots.
tests/test_okls.py	Adds CUDA optimizer smoke, state-initialization, and dimensionality-validation coverage.
tests/test_matrix_root_inverse_utils.py	Adds shape, precision-restoration, accuracy, and dtype-validation coverage for inverse roots.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    G[Gradient] --> I[Initialize or update Kronecker factors]
    I --> C[Scaled CANS inverse roots]
    G --> M[Nesterov momentum]
    C --> P[Two-sided preconditioning]
    M --> P
    P --> U[Parameter update]

Loading

_{Reviews (2): Last reviewed commit: "fix: use decoupled weight decay in OKLS" | Re-trigger Greptile}

[xsgxlz]

Thanks for upstreaming OKLS! I opened a follow-up PR against your branch: mkhona-nvidia/Emerging-Optimizers#2.

It adds:

• mixed, high, and highest scaled-CANS precision modes; mixed uses FP16 GEMMs with FP32 accumulation to avoid the NaNs observed with BF16 CANS, aligning with the Scaled CANS NS we proposed in the blog
• Packed upper-triangular storage for symmetric Kronecker factors and inverse roots, reducing persistent symmetric-state memory by 50%

The changes passed 24 CUDA tests, Ruff checks, and multi-seed MLP trajectory comparisons against both the public release and internal implementation.

[skyw]

Thanks for upstreaming OKLS! I opened a follow-up PR against your branch: mkhona-nvidia/Emerging-Optimizers#2.

It adds:

• mixed, high, and highest scaled-CANS precision modes; mixed uses FP16 GEMMs with FP32 accumulation to avoid the NaNs observed with BF16 CANS, aligning with the Scaled CANS NS we proposed in the blog
• Packed upper-triangular storage for symmetric Kronecker factors and inverse roots, reducing persistent symmetric-state memory by 50%

The changes passed 24 CUDA tests, Ruff checks, and multi-seed MLP trajectory comparisons against both the public release and internal implementation.

Thanks for the contribution. Our precision supports was mostly following what Pytorch would do for fp32 input. The bf16 cast was mainly because of pytorch lack of "medium" kernels. It would be a very significant testing cost if we expand beyond that and don't have plan to do so.

Packed storage is a great idea. Implementation will need more consideration as we'll need to incorporate rest of the code base as well as other NVIDIA SW stack's (megatron for example) requirement, which is out of this PR's scope.

[skyw]

Also this PR was for demonstration and test.
We are changing abstraction of our shampoo/soap implementation (e.g. WIP https://github.com/NVIDIA-NeMo/Emerging-Optimizers/blob/skyw/okls_exp/emerging_optimizers/experimental/okls.py), so specific combination can be more contained and hopefully we have more room to accept things like fp16 support for just one version of OKLS.

[xsgxlz]

Also this PR was for demonstration and test. We are changing abstraction of our shampoo/soap implementation (e.g. WIP https://github.com/NVIDIA-NeMo/Emerging-Optimizers/blob/skyw/okls_exp/emerging_optimizers/experimental/okls.py), so specific combination can be more contained and hopefully we have more room to accept things like fp16 support for just one version of OKLS.

Thanks for the clarification. Because we found that, unlike Muon, bf16 almost certainly leads to a NaN for OKLS training, do you think we should remove the "medium" option? Or raise a warning for it?

[xsgxlz]

Besides, we found that a tight estimate for the largest eigenvalue is helpful, especially when the model is large. Please consider this upper bound in our repo if it aligns with your requirements

def _estimate_max_eigenvalue(A: torch.Tensor) -> torch.Tensor:
    """Strict upper bound via min(Wolkowicz-Styan, Minc-Sainte-Marie). Cost: O(n²)."""
    n = A.size(-1)
    diag = A.diagonal(dim1=-2, dim2=-1)

    # ── Optimal scaling: max(|A/c|) = √FP32_MAX / n ──
    # Both bounds involve O(n²·max²) intermediates.  Scaling by
    # c = max(|diag|)·n/√FP32_MAX keeps every intermediate in FP32 range
    # while maximising dynamic-range usage (fewest small elements lost).
    # For SPD matrices max(|A_ij|) ≤ max(diag_i), so this is a safe bound.
    _FROB_SCALE = n / math.sqrt(torch.finfo(torch.float32).max)  # n / 1.844e19
    c = (diag.abs().max(dim=-1).values * _FROB_SCALE).clamp(min=1e-30)
    c_inv = 1.0 / c

    # Single n² materialisation: |A/c|  (||x||_F = |||x|||_F, so both bounds
    # can work from the elementwise-abs form without loss of information.)
    abs_A_s = (A * c_inv[..., None, None]).abs()

    # ── Wolkowicz-Styan bound on A/c, then unscale ──
    m_s = diag.sum(dim=-1) * (c_inv / n)
    f_s_n = torch.linalg.matrix_norm(abs_A_s) * (1.0 / math.sqrt(n))
    s_s = torch.sqrt(torch.clamp((f_s_n + m_s) * (f_s_n - m_s), min=0.0))
    ws_bound = c * (m_s + s_s * math.sqrt(n - 1))

    # ── Minc-Sainte-Marie bound on A/c, then unscale ──
    d_s = torch.sum(abs_A_s, dim=-1)
    d_s_clamped = torch.clamp(d_s, min=1e-12)
    y_s = torch.einsum("...ij,...j->...i", abs_A_s, d_s_clamped)
    minc_bound = c * torch.max(y_s / d_s_clamped, dim=-1).values

    return torch.minimum(ws_bound, minc_bound)

[skyw]

Also this PR was for demonstration and test. We are changing abstraction of our shampoo/soap implementation (e.g. WIP https://github.com/NVIDIA-NeMo/Emerging-Optimizers/blob/skyw/okls_exp/emerging_optimizers/experimental/okls.py), so specific combination can be more contained and hopefully we have more room to accept things like fp16 support for just one version of OKLS.

Thanks for the clarification. Because we found that, unlike Muon, bf16 almost certainly leads to a NaN for OKLS training, do you think we should remove the "medium" option? Or raise a warning for it?

Do you have a test case that CANS straight blow up (inf or nan) with bf16? That would be a sufficient reason to disallow bf16.
If it indirectly gets NaN deep in a training pipeline, probably default to TF32 and pop a warning for lower.

[xsgxlz]

Also this PR was for demonstration and test. We are changing abstraction of our shampoo/soap implementation (e.g. WIP https://github.com/NVIDIA-NeMo/Emerging-Optimizers/blob/skyw/okls_exp/emerging_optimizers/experimental/okls.py), so specific combination can be more contained and hopefully we have more room to accept things like fp16 support for just one version of OKLS.

Thanks for the clarification. Because we found that, unlike Muon, bf16 almost certainly leads to a NaN for OKLS training, do you think we should remove the "medium" option? Or raise a warning for it?

Do you have a test case that CANS straight blow up (inf or nan) with bf16? That would be a sufficient reason to disallow bf16. If it indirectly gets NaN deep in a training pipeline, probably default to TF32 and pop a warning for lower.

Yes - with BF16, the NS will almost surely explode for moderately large or ill-conditioned matrices. In our real training, we found BF16 survives at most 3 steps, so it might be better to disallow it.

import torch

from emerging_optimizers.soap.matrix_root_inverse_utils import (
    mat_root_inv_via_scaled_cans,
)

device = "cuda"
size = 1024

generator = torch.Generator(device=device).manual_seed(0)
x = torch.randn(size, size, generator=generator, device=device)

matrix = (x @ x.T) / size
matrix.diagonal().add_(1e-6)

with torch.no_grad():
    for precision in ("medium", "high"):
        inverse_root = mat_root_inv_via_scaled_cans(
            matrix,
            fp32_matmul_prec=precision,
        )

        print(f"{precision}:")
        print("  finite:", torch.isfinite(inverse_root).sum().item())
        print("  NaN:", torch.isnan(inverse_root).sum().item())
        print("  Inf:", torch.isinf(inverse_root).sum().item())

medium:
  finite: 0
  NaN: 413232
  Inf: 635344
high:
  finite: 1048576
  NaN: 0
  Inf: 0