Sync 3 diffs to facebookresearch/optimizers (#276)

CHANGELOG.md

@@ -1,13 +1,20 @@
 # Changelog
 
-Release notes for the OSS mirror, generated by `scripts/oss_sync.sh` from the internal source tree.
+Release notes for the OSS mirror. Seeded by `scripts/oss_sync.sh` from the internal source tree and edited for release.
+
+## 2026-08-19
+
+- [shampoo] Add `NewtonSchulzRootInvConfig`: a matmul-only coupled Newton-Schulz iteration, usable as the `amortized_computation_config` of `RootInvShampooPreconditionerConfig` in place of the default eigendecomposition. Opt-in — `DefaultShampooConfig` is unchanged. The iteration runs a fixed number of steps with no residual-based stopping criterion, so it introduces no host-device synchronization. Only power-of-two inverse roots are supported; anything else raises at optimizer construction rather than degrading to a per-factor-matrix warning that reuses a stale preconditioner. `relative_epsilon` (default `1e-6`) floors the ridge at a fraction of `|A|_F` and is applied unconditionally, so on the rank-deficient factor matrices seen early in training this path regularizes more aggressively than eigendecomposition at the same epsilon and the two do not agree there. `coefficients` takes a per-iteration schedule of `(a, b, c)` triples for `p(x) = a x + b x^3 + c x^5`, defaulting to a 10-step Polar Express schedule; tf32 is disabled inside the iteration by default (`disable_tf32=True`).
+- [shampoo] Annotate `ShampooPT2CompileConfig`'s `make_dataclass`-synthesized base class `Any` so mypy accepts it as a base class, and add test coverage asserting the synthesized field set matches `torch.compile`'s signature (no runtime change).
+- [shampoo] Remove unused type-error suppression comments across the optimizer, examples, and tests (no behavior change).
 
 ## 2026-07-24
 
 - [shampoo] Drop the experimental underscore prefix and note from the now-productionized FullyShard/HybridShard lossless distributor modules (rename only, no behavior change).
 - [shampoo] Balance per-owner-rank byte load in the FSDP/HSDP lossless ROUND_ROBIN distributor and correctly support `num_sub_groups > 1` via per-bucket LPT owner assignment.
 - [shampoo] Store the per-group `step`, `lr_sum`, and `train_mode` scalars under every parameter's state so the checkpoint is param-keyed and DCP resume survives `num_sub_groups` changes.
 - [Shampoo] Fix an FSDP2/HSDP2 lossless distributor crash on the first step after checkpoint resume when a parameter receives no gradient, by rebuilding masked lists consistently.
+- [shampoo] Fix a mypy failure in the test suite by assigning the `Tensor` cast to a new local instead of shadowing the `Parameter` loop variable (no behavior change).
 
 ## 2026-07-20
 

distributed_shampoo/__init__.py

@@ -23,6 +23,7 @@
     EighEigendecompositionConfig,
     MatrixFunctionConfig,
     NewtonSchulzOrthogonalizationConfig,
+    NewtonSchulzRootInvConfig,
     OrthogonalizationConfig,
     PerturbationConfig,
     PseudoInverseConfig,
@@ -118,6 +119,7 @@
     "DefaultEigenConfig",  # Default `RootInvConfig` using `EigenConfig`.
     "CoupledNewtonConfig",  # Based on `RootInvConfig`.
     "CoupledHigherOrderConfig",  # Based on `RootInvConfig`.
+    "NewtonSchulzRootInvConfig",  # Based on `RootInvConfig`.
     "OrthogonalizationConfig",  # Abstract base class (based on `MatrixFunctionConfig`).
     "SVDOrthogonalizationConfig",  # Based on `OrthogonalizationConfig`.
     "NewtonSchulzOrthogonalizationConfig",  # Based on `OrthogonalizationConfig`.

distributed_shampoo/distributed_shampoo.py

@@ -981,7 +981,6 @@ def _instantiate_per_group_step(
         # Use PT2 to compile the step function for each parameter group.
         self._per_group_step: Callable[..., None] = (
             torch.compile(
-                # pyrefly: ignore [bad-argument-type]
                 self._per_group_step_impl,
                 # pyrefly: ignore [bad-argument-type]
                 **asdict(shampoo_pt2_compile_config),
@@ -1370,7 +1369,7 @@ def _apply_in_place_primal_averaging(
         # This computes: - (1 - mu_x * mu_y) * lr * P.
         torch._foreach_mul_(
             masked_blocked_search_directions,
-            (1 - train_interp_coeff * eval_interp_coeff),  # type: ignore
+            (1 - train_interp_coeff * eval_interp_coeff),
         )
         # This computes: (1 - mu_x) * (Z_old - Y) - (1 - mu_x * mu_y) * lr * P.
         # pyrefly: ignore [no-matching-overload]

distributed_shampoo/examples/utils.py

@@ -73,7 +73,6 @@ def get_distributed_env() -> tuple[int, int, int]:
 
 def set_seed(seed: int) -> None:
     torch.manual_seed(seed)
-    # pyrefly: ignore [bad-argument-type]
     np.random.seed(seed)
     random.seed(seed)
     torch.use_deterministic_algorithms(True)

distributed_shampoo/preconditioner/README.md

@@ -330,6 +330,7 @@ def matrix_inverse_root(
 - **Eigendecomposition**: Most stable, best for symmetric positive definite matrices
 - **Newton Iteration**: Fast convergence for well-conditioned matrices
 - **Higher-Order Coupled**: Advanced methods for fractional powers
+- **Newton-Schulz** (`NewtonSchulzRootInvConfig`): opt-in, matmul-only; only power-of-two roots (orders 1, 2, 4)
 
 #### Eigendecomposition
 ```python

distributed_shampoo/preconditioner/matrix_functions.py

@@ -31,6 +31,7 @@
     EigendecompositionConfig,
     EighEigendecompositionConfig,
     NewtonSchulzOrthogonalizationConfig,
+    NewtonSchulzRootInvConfig,
     OrthogonalizationConfig,
     PerturbationConfig,
     PseudoInverseConfig,
@@ -436,6 +437,128 @@ def matrix_inverse_root_newton(
 
         return X, M, termination_flag, iteration, error
 
+    def matrix_inverse_root_newton_schulz(
+        A: Tensor,
+        root: int,
+        coefficients: list[list[float]],
+        epsilon: float = 0.0,
+        relative_epsilon: float = 1e-6,
+        disable_tf32: bool = True,
+    ) -> Tensor:
+        """Compute matrix inverse root using the coupled Newton-Schulz iteration.
+
+        A single pass implements SqrtInverseNewtonSchulz, which returns both A^{1/2} and A^{-1/2}:
+
+            alpha <- |A|_F
+            Y <- A / alpha, Z <- I
+            for each (a, b, c) in the coefficient schedule
+                T <- Z Y
+                B <- b T + c T^2
+                Y <- a Y + Y B
+                Z <- a Z + B Z
+            A^{1/2} ~= sqrt(alpha) Y, A^{-1/2} ~= Z / sqrt(alpha)
+
+        Y and Z are updated as a coupled pair rather than by recomputing a residual from A, which is
+        what makes the iteration numerically stable; the uncoupled form diverges within ~15 iterations
+        even in double precision.
+
+        Only roots that are powers of two are supported. Since a pass yields both the square root and
+        the inverse square root, A^{-1/2^m} is obtained by chaining m passes: each of the first m - 1
+        passes feeds its square root output forward, and the last pass returns its inverse square root
+        output.
+
+        NOTE: Coefficients summing to 1 give the scalar map a fixed point at 1, so the iteration
+            converges to the inverse square root. Coefficients that do not sum to 1 -- such as the Muon
+            coefficients (3.4445, -4.7750, 2.0315) used by newton_schulz for orthogonalization, where
+            only the sign of the singular values matters -- instead converge to a band around it.
+            NewtonSchulzRootInvConfig warns about this at construction time.
+
+        NOTE: Unlike the eigendecomposition path, this iteration cannot stabilize a rank-deficient or
+            indefinite input, because it never forms the spectrum it would need to shift. Shampoo
+            factor matrices are rank-deficient early in training and pick up small negative eigenvalues
+            from floating point error, on which Z diverges. relative_epsilon guards against this by
+            flooring the ridge at a fraction of |A|_F, which dominates those spurious eigenvalues and
+            bounds the condition number. The consequence is that on rank-deficient input this function
+            regularizes more aggressively than the eigendecomposition path does for the same epsilon,
+            and therefore does not agree with it there.
+
+        References:
+            - https://arxiv.org/abs/2505.16932 (Polar Express)
+            - https://docs.modula.systems/algorithms/newton-schulz/
+
+        Args:
+            A (Tensor): Matrix of interest.
+            root (int): Root of interest. Must be a power of two.
+            coefficients (list[list[float]]): Per-iteration schedule of (a, b, c) coefficient
+                triples for the odd polynomial p(x) = a x + b x^3 + c x^5 driving the iteration. The
+                number of iterations is len(coefficients). Validated by
+                NewtonSchulzRootInvConfig.__post_init__, which is the only supported way to reach
+                this function.
+            epsilon (float): Adds epsilon * I to matrix before taking matrix root. (Default: 0.0)
+            relative_epsilon (float): Adds relative_epsilon * |A|_F * I to the matrix before taking the
+                matrix root, taking the larger of this and epsilon. Required for rank-deficient input;
+                see the note above. (Default: 1e-6)
+            disable_tf32 (bool): Whether to disable tf32 matmuls or not internally. Highly recommend
+                keeping True; tf32 cannot represent the relative_epsilon-scale eigenvalues this
+                iteration must resolve, and Z @ Y is not a Gram product, so the resulting error is
+                unstructured and diverges. (Default: True)
+
+        Returns:
+            X (Tensor): Inverse root of matrix A.
+
+        Raises:
+            ValueError: If root is not a power of two.
+
+        """
+        # This should not be reachable: RootInvShampooPreconditionerList already rejects roots that
+        # are not powers of two at optimizer construction time. Kept as a defensive guard in case
+        # this function is ever called directly, bypassing that validation.
+        if root < 2 or root & (root - 1):
+            raise ValueError(
+                f"{root=} must be a power of two to use Newton-Schulz iteration!"
+            )
+
+        tf32_flag = torch.backends.cuda.matmul.allow_tf32
+        if disable_tf32:
+            torch.backends.cuda.matmul.allow_tf32 = False
+
+        try:
+            # Add regularization, floored relative to |A|_F so that the input is positive definite. The
+            # ridge is kept as a 0-d tensor rather than a float so that no host-device synchronization
+            # is introduced.
+            identity = torch.eye(A.shape[0], dtype=A.dtype, device=A.device)
+            A = torch.addcmul(
+                A,
+                identity,
+                torch.linalg.matrix_norm(A).mul_(relative_epsilon).clamp_min_(epsilon),
+            )
+
+            # A^{-1/root} = ((...(A^{1/2})^{1/2}...)^{1/2})^{-1/2} with log2(root) nested square roots.
+            for pass_index in range(root.bit_length() - 1, 0, -1):
+                # Normalize so that the spectrum of Y lies in (0, 1]; |A|_2 <= |A|_F.
+                alpha = torch.linalg.matrix_norm(A).clamp_min(1e-8)
+                Y = A / alpha
+                # Cloned because the final pass returns Z.div_(sqrt_alpha), which mutates in place,
+                # and identity is reused across passes. torch.addmm below is out-of-place and would
+                # rebind Z on its own, but that only holds while the schedule is non-empty.
+                Z = identity.clone()
+
+                for a, b, c in coefficients:
+                    T = Z @ Y
+                    B = torch.addmm(T, T, T, beta=b, alpha=c)
+                    Y = torch.addmm(Y, Y, B, beta=a, alpha=1)
+                    Z = torch.addmm(Z, B, Z, beta=a, alpha=1)
+
+                sqrt_alpha = alpha.sqrt()
+                # Feed A^{1/2} into the next pass; the final pass produces the inverse root.
+                A = Y.mul_(sqrt_alpha) if pass_index > 1 else Z.div_(sqrt_alpha)
+        finally:
+            # Always restore tf32 mode unconditionally, so we skip the disable_tf32 check. When
+            # disable_tf32=False, this is a no-op since tf32_flag already equals the current value.
+            torch.backends.cuda.matmul.allow_tf32 = tf32_flag
+
+        return A
+
     def matrix_inverse_root_higher_order(
         A: Tensor,
         root: Fraction,
@@ -679,6 +802,16 @@ def matrix_inverse_root_higher_order(
                 logger.warning(
                     "Newton did not converge and reached maximum number of iterations!"
                 )
+        case NewtonSchulzRootInvConfig():
+            # NOTE: Use Fraction.is_integer() instead when downstream applications are Python 3.12+ available
+            if root.denominator != 1:
+                raise ValueError(
+                    f"{root.denominator=} must be equal to 1 to use Newton-Schulz iteration!"
+                )
+
+            X = _assign_function_args_from_config(
+                func=matrix_inverse_root_newton_schulz, config=root_inv_config
+            )(A=A, root=root.numerator, epsilon=epsilon)
         case CoupledHigherOrderConfig():
             X, _, termination_flag, _, _ = _assign_function_args_from_config(
                 func=matrix_inverse_root_higher_order, config=root_inv_config

distributed_shampoo/preconditioner/matrix_functions_types.py

@@ -7,11 +7,15 @@
 
 """
 
+import logging
+import math
 from collections.abc import Callable
 from dataclasses import dataclass, field
 
 from distributed_shampoo.utils.abstract_dataclass import AbstractDataclass
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 @dataclass(init=False)
 class RankDeficientStabilityConfig(AbstractDataclass):
@@ -213,6 +217,86 @@ class CoupledNewtonConfig(RootInvConfig):
     tolerance: float = 1e-6
 
 
+@dataclass(kw_only=True)
+class NewtonSchulzRootInvConfig(RootInvConfig):
+    """Configuration for matrix root inverse via the coupled Newton-Schulz iteration.
+
+    Unlike CoupledNewtonConfig and CoupledHigherOrderConfig, the iteration runs for a fixed number of
+    steps and never evaluates a residual-based stopping criterion, so it incurs no host-device
+    synchronization and is expressed entirely as matmuls. Only roots that are powers of two are
+    supported, i.e. blocks of order 1, 2, and 4.
+
+    WARNING: On rank-deficient input this regularizes more aggressively than the eigendecomposition
+        path does for the same epsilon, so the two do not agree there. See relative_epsilon.
+
+    Attributes:
+        relative_epsilon (float): Floors the ridge added before the iteration at
+            relative_epsilon * |A|_F. Unlike the eigendecomposition path, this iteration cannot
+            stabilize a rank-deficient or indefinite matrix, and Shampoo factor matrices are
+            rank-deficient early in training; without this floor the iteration diverges to NaN.
+            (Default: 1e-6)
+        coefficients (list[list[float]]): Per-iteration schedule of (a, b, c) coefficient triples
+            for the odd polynomial p(x) = a x + b x^3 + c x^5 driving the iteration. The number
+            of iterations is len(coefficients).
+            (Default: Polar Express 10-step schedule from ASGO.)
+        disable_tf32 (bool): Whether to disable tf32 matmuls or not internally. Highly recommend
+            keeping True. The iteration is built entirely out of matmuls and must resolve
+            eigenvalues down to relative_epsilon, which tf32's 10-bit mantissa cannot represent;
+            unlike the factor matrix accumulation, Z @ Y is not a Gram product, so tf32 error there
+            is unstructured and drives the iteration to NaN. (Default: True)
+
+    """
+
+    @staticmethod
+    def _get_default_coefficients() -> list[list[float]]:
+        return [
+            [8.28721201814563, -23.595886519098837, 17.300387312530933],
+            [4.107059111542203, -2.9478499167379106, 0.5448431082926601],
+            [3.9486908534822946, -2.9089021159629490, 0.5518191394370137],
+            [3.3184196573706015, -2.4884880243148740, 0.5100489401237200],
+            [2.300652019954817, -1.6689039845747493, 0.4188073119525673],
+            [1.891301407787398, -1.2679958271945868, 0.3768040894852483],
+            [1.8750014808534479, -1.2500016453999487, 0.3750001645474248],
+            [1.875, -1.25, 0.375],
+            [1.875, -1.25, 0.375],
+            [1.875, -1.25, 0.375],
+        ]
+
+    relative_epsilon: float = 1e-6
+    disable_tf32: bool = True
+    # TODO: Clean up coefficient definition -- consider using list[tuple[float, float, float]]
+    # to enforce 3-tuples, and define defaults from the training pipeline side.
+    coefficients: list[list[float]] = field(default_factory=_get_default_coefficients)
+
+    def __post_init__(self) -> None:
+        if len(self.coefficients) == 0:
+            raise ValueError("coefficients must be non-empty.")
+        for index, entry in enumerate(self.coefficients):
+            if len(entry) != 3:
+                raise ValueError(
+                    f"coefficients[{index}] must contain exactly three coefficients (a, b, c) for "
+                    f"p(x) = a x + b x^3 + c x^5, but {entry=} has {len(entry)}."
+                )
+            for coefficient in entry:
+                if isinstance(coefficient, bool) or not isinstance(
+                    coefficient, (int, float)
+                ):
+                    raise ValueError(
+                        f"coefficients[{index}] must contain real numbers, but {entry=} contains "
+                        f"{coefficient!r} of type {type(coefficient).__name__}."
+                    )
+                if not math.isfinite(coefficient):
+                    raise ValueError(
+                        f"coefficients[{index}] must be finite, but {entry=} contains {coefficient}."
+                    )
+        final_coefficients = self.coefficients[-1]
+        if not math.isclose(sum(final_coefficients), 1.0):
+            logger.warning(
+                f"{final_coefficients=} do not sum to 1, so the Newton-Schulz iteration has no fixed "
+                "point at 1 and will converge to a band around the inverse root rather than to it."
+            )
+
+
 @dataclass(kw_only=True)
 class CoupledHigherOrderConfig(RootInvConfig):
     """Configuration for matrix root inverse via coupled higher-order method.

distributed_shampoo/preconditioner/shampoo_preconditioner_list.py

@@ -29,6 +29,7 @@
 from distributed_shampoo.preconditioner.matrix_functions_types import (
     EigendecompositionConfig,
     MatrixFunctionConfig,
+    NewtonSchulzRootInvConfig,
     RootInvConfig,
 )
 from distributed_shampoo.preconditioner.preconditioner_list import (
@@ -686,6 +687,26 @@ def __post_init__(self) -> None:
             == len(self.factor_matrices)
             == len(self.inv_factor_matrices)
         )
+        # Fail fast rather than at the first amortized computation: an input the amortized
+        # computation cannot handle would otherwise surface as a swallowed per-factor-matrix warning
+        # that silently reuses the stale preconditioner until
+        # num_tolerated_failed_amortized_computations is hit.
+        if isinstance(self.amortized_computation_config, NewtonSchulzRootInvConfig):
+            if unsupported_roots := sorted(
+                {
+                    root
+                    for root in self.roots
+                    if not float(root).is_integer()
+                    or int(root) < 2
+                    or int(root) & (int(root) - 1)
+                }
+            ):
+                raise ValueError(
+                    f"{type(self.amortized_computation_config).__name__} only supports inverse roots "
+                    f"that are powers of two, but {unsupported_roots=} were requested. Merge or block "
+                    "the offending parameters down to order 1, 2, or 4, or set inverse_exponent_override "
+                    "to the reciprocal of a power of two."
+                )
 
 
 @dataclass(kw_only=True)