Add klshampoo (#302)

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

Author: skyw

emerging_optimizers/shampoo/__init__.py

@@ -13,5 +13,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from emerging_optimizers.shampoo.shampoo import Shampoo
+from emerging_optimizers.shampoo.shampoo import KlShampoo, Shampoo
 from emerging_optimizers.shampoo.soap_v3 import KlMSoap, KlSoapV3, ReklsV3

emerging_optimizers/shampoo/shampoo.py

@@ -12,6 +12,7 @@
 # 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, ClassVar, override
 
 
@@ -23,18 +24,36 @@
 from torch.optim.optimizer import ParamsT
 
 from emerging_optimizers import mixin as opt_mixin
-from emerging_optimizers import registry
+from emerging_optimizers import registry, utils
+from emerging_optimizers.legacy_soap import soap
 from emerging_optimizers.shampoo import precond_base
 from emerging_optimizers.utils import eig as eig_utils
 
 
 __all__ = [
+    "KlShampoo",
+    "KlShampooPreconditioner",
     "Shampoo",
     "ShampooBase",
     "ShampooPreconditioner",
 ]
 
 
+def _get_root_inverse_from_eigens(
+    eigvals: torch.Tensor,
+    eigvecs: torch.Tensor,
+    p_root_inv: float,
+    eps: float,
+) -> torch.Tensor:
+    # Eigh can sometime return negative values for numerical 0; clamping to 0 removes them
+    eigvals = eigvals.clamp_min(0)
+
+    # Tikhonov regularization
+    exp = 1.0 / p_root_inv
+    inv_root_scale = eigvals**exp / (eigvals ** (2 * exp) + eps ** (2 * exp))
+    return (eigvecs * inv_root_scale) @ eigvecs.mT
+
+
 class ShampooPreconditioner:
     """Per-parameter Shampoo preconditioner holding the Kronecker factors of one 2D parameter.
 
@@ -128,38 +147,106 @@ def step(self, grad: torch.Tensor, shampoo_beta: float) -> None:
         """
         self.update_kronecker_factors(grad, shampoo_beta)
 
-    def _get_root_inverse(self, kronecker_factor: torch.Tensor) -> torch.Tensor:
-        """Computes ``kronecker_factor^(-1/p_root_inv)`` from its eigendecomposition.
+    def precondition(self, x: torch.Tensor) -> torch.Tensor:
+        """Applies both root inverse to a matrix in the parameter basis.
 
         Args:
-            kronecker_factor: left or right kronecker factor
+            x: Matrix in the parameter basis.
 
         Returns:
-            The inverse root of the factor.
+            The preconditioned matrix, in the parameter basis.
         """
-        eigvals, eigvecs = eig_utils.eigh_with_fallback(kronecker_factor)
+        eigvals_L, eigvecs_L = eig_utils.eigh_with_fallback(self.kronecker_factor_pair.L)
+        eigvals_R, eigvecs_R = eig_utils.eigh_with_fallback(self.kronecker_factor_pair.R)
 
-        # Eigh can sometime return negative values for numerical 0; clamping to 0 removes them
-        eigvals = eigvals.clamp_min(0)
+        root_inv_L = _get_root_inverse_from_eigens(eigvals_L, eigvecs_L, self.p_root_inv, self.eps)
+        root_inv_R = _get_root_inverse_from_eigens(eigvals_R, eigvecs_R, self.p_root_inv, self.eps)
 
-        # Tikhonov regularization
-        exp = 1.0 / self.p_root_inv
-        inv_root_scale = eigvals**exp / (eigvals ** (2 * exp) + self.eps ** (2 * exp))
-        return (eigvecs * inv_root_scale) @ eigvecs.mT
+        return root_inv_L @ x @ root_inv_R
 
-    def precondition(self, x: torch.Tensor) -> torch.Tensor:
-        """Applies both root inverse to a matrix in the parameter basis.
 
-        Args:
-            x: Matrix in the parameter basis.
+class KlShampooPreconditioner(ShampooPreconditioner):
+    """Shampoo preconditioner with the KL-corrected Kronecker factor update."""
 
-        Returns:
-            The preconditioned matrix, in the parameter basis.
+    def __init__(
+        self,
+        state: dict,
+        p_root_inv: float,
+        eps: float,
+    ) -> None:
+        super().__init__(state, p_root_inv, eps)
+        self.eigenbasis_pair = precond_base.TensorPair(state["Q_L"], state["Q_R"])
+        self.eigvals_pair = precond_base.TensorPair(state["eigvals_L"], state["eigvals_R"])
+
+    @staticmethod
+    @override
+    def init_state(
+        shape: tuple[int, ...],
+        device: torch.device,
+    ) -> dict[str, torch.Tensor]:
+        state = ShampooPreconditioner.init_state(shape, device)
+        m, n = shape
+        state.update(
+            {
+                "Q_L": torch.eye(m, device=device),
+                "Q_R": torch.eye(n, device=device),
+                "eigvals_L": torch.zeros(m, device=device),
+                "eigvals_R": torch.zeros(n, device=device),
+            }
+        )
+        return state
+
+    @override
+    def rebind_state(self, state: dict) -> None:
+        updates = {
+            "L": self.kronecker_factor_pair.L,
+            "R": self.kronecker_factor_pair.R,
+            "Q_L": self.eigenbasis_pair.L,
+            "Q_R": self.eigenbasis_pair.R,
+            "eigvals_L": self.eigvals_pair.L,
+            "eigvals_R": self.eigvals_pair.R,
+        }
+        missing = updates.keys() - state.keys()
+        if missing:
+            raise KeyError(f"rebind_state: state missing keys {sorted(missing)}")
+        state.update(updates)
+
+    @override
+    def init_step(self, grad: torch.Tensor, shampoo_beta: float) -> None:
+        """Performs the first step's factor update, before any history exists.
+
+        Args:
+            grad: Gradient of the parameter on the first step.
+            shampoo_beta: EMA coefficient for the Kronecker factor update.
         """
-        root_inv_L = self._get_root_inverse(self.kronecker_factor_pair.L)
-        root_inv_R = self._get_root_inverse(self.kronecker_factor_pair.R)
+        self.update_kronecker_factors(grad, shampoo_beta)
 
-        return root_inv_L @ x @ root_inv_R
+    @override
+    def update_kronecker_factors(self, grad: torch.Tensor, shampoo_beta: float) -> None:
+        with utils.fp32_matmul_precision("highest"):
+            soap.update_kronecker_factors_kl_shampoo(
+                self.kronecker_factor_pair,
+                grad,
+                shampoo_beta,
+                self.eigenbasis_pair,
+                self.eigvals_pair,
+                self.eps,
+            )
+
+    @override
+    def precondition(self, x: torch.Tensor) -> torch.Tensor:
+        eigvals_L, eigvecs_L = eig_utils.eigh_with_fallback(self.kronecker_factor_pair.L)
+        eigvals_R, eigvecs_R = eig_utils.eigh_with_fallback(self.kronecker_factor_pair.R)
+
+        self.eigenbasis_pair = precond_base.TensorPair(eigvecs_L, eigvecs_R)
+        self.eigvals_pair = precond_base.TensorPair(eigvals_L, eigvals_R)
+
+        root_inv_L = _get_root_inverse_from_eigens(eigvals_L, eigvecs_L, self.p_root_inv, self.eps)
+        root_inv_R = _get_root_inverse_from_eigens(eigvals_R, eigvecs_R, self.p_root_inv, self.eps)
+
+        m, n = x.shape
+        shape_scale = math.sqrt(m / n) / (math.sqrt(m) + math.sqrt(n))
+        return (root_inv_L @ x @ root_inv_R) * shape_scale
 
 
 class ShampooBase(optim.Optimizer, opt_mixin.WeightDecayMixin):
@@ -286,7 +373,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
         Args:
             closure: Unsupported; must be ``None``.
 
-        Raises:f
+        Raises:
             ValueError: If ``closure`` is not ``None``.
         """
         if closure is not None:
@@ -359,3 +446,31 @@ def _scalar_update(
         """
         exp_avg.lerp_(grad, 1 - momentum)
         return exp_avg
+
+
+@registry.register_optimizer("kl_shampoo")
+class KlShampoo(Shampoo):
+    """Shampoo with KL-corrected Kronecker factors and EMA momentum."""
+
+    PreconditionerCls: ClassVar[type[precond_base.ShampooPreconditionerProtocol]] = KlShampooPreconditioner
+
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: float,
+        momentum: float = 0.9,
+        shampoo_beta: float = 0.95,
+        eps: float = 1e-8,
+        weight_decay: float = 0.01,
+        *,
+        p_root_inv: float = 4,
+    ) -> None:
+        super().__init__(
+            params,
+            lr,
+            momentum,
+            shampoo_beta,
+            eps,
+            weight_decay,
+            p_root_inv=p_root_inv,
+        )