diff --git a/CHANGELOG.md b/CHANGELOG.md index 31e8e92..4a1895a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,12 @@ # 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 @@ -8,6 +14,7 @@ Release notes for the OSS mirror, generated by `scripts/oss_sync.sh` from the in - [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 diff --git a/distributed_shampoo/__init__.py b/distributed_shampoo/__init__.py index 767cf6a..2ec751d 100644 --- a/distributed_shampoo/__init__.py +++ b/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`. diff --git a/distributed_shampoo/distributed_shampoo.py b/distributed_shampoo/distributed_shampoo.py index 0bbbc84..bdf8014 100644 --- a/distributed_shampoo/distributed_shampoo.py +++ b/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] diff --git a/distributed_shampoo/examples/utils.py b/distributed_shampoo/examples/utils.py index a237c93..291a294 100644 --- a/distributed_shampoo/examples/utils.py +++ b/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) diff --git a/distributed_shampoo/preconditioner/README.md b/distributed_shampoo/preconditioner/README.md index de5dbf0..bad331c 100644 --- a/distributed_shampoo/preconditioner/README.md +++ b/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 diff --git a/distributed_shampoo/preconditioner/matrix_functions.py b/distributed_shampoo/preconditioner/matrix_functions.py index 48347e4..a67c8de 100644 --- a/distributed_shampoo/preconditioner/matrix_functions.py +++ b/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 diff --git a/distributed_shampoo/preconditioner/matrix_functions_types.py b/distributed_shampoo/preconditioner/matrix_functions_types.py index 02e7c63..ade8cf3 100644 --- a/distributed_shampoo/preconditioner/matrix_functions_types.py +++ b/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. diff --git a/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py b/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py index d4a6732..db57717 100644 --- a/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py +++ b/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) diff --git a/distributed_shampoo/preconditioner/tests/matrix_functions_test.py b/distributed_shampoo/preconditioner/tests/matrix_functions_test.py index 0ccf5b7..296e051 100644 --- a/distributed_shampoo/preconditioner/tests/matrix_functions_test.py +++ b/distributed_shampoo/preconditioner/tests/matrix_functions_test.py @@ -8,6 +8,7 @@ """ import itertools +import math import re import unittest from collections.abc import Callable @@ -37,6 +38,7 @@ EigendecompositionConfig, EighEigendecompositionConfig, NewtonSchulzOrthogonalizationConfig, + NewtonSchulzRootInvConfig, OrthogonalizationConfig, PerturbationConfig, PseudoInverseConfig, @@ -661,6 +663,207 @@ def A_tridiagonal_2(n: int, alpha: float, beta: float) -> Tensor: ) +@instantiate_parametrized_tests +class NewtonSchulzRootInverseTest(unittest.TestCase): + @staticmethod + def _spd_matrix(n: int, condition_number: float) -> Tensor: + torch.manual_seed(42) + Q, _ = torch.linalg.qr(torch.randn(n, n, dtype=torch.float64)) + eigenvalues = torch.logspace( + 0, -math.log10(condition_number), n, dtype=torch.float64 + ) + return ((Q * eigenvalues) @ Q.T).float() + + @staticmethod + def _relative_error(X: Tensor, expected: Tensor) -> float: + """Normwise relative error. Elementwise rtol is not meaningful here: entries of the inverse + root that are near zero carry large relative error at negligible absolute error.""" + return ( + torch.dist(X, expected, p=torch.inf) + / torch.linalg.norm(expected, ord=torch.inf) + ).item() + + @parametrize("root", [2, 4, 8]) + @parametrize("n", [10, 100]) + def test_newton_schulz_root_inverse_identity(self, n: int, root: int) -> None: + torch.testing.assert_close( + matrix_inverse_root( + A=torch.eye(n), + root=Fraction(root), + root_inv_config=NewtonSchulzRootInvConfig(), + ), + torch.eye(n), + atol=1e-5, + rtol=1e-5, + ) + + @parametrize("root", [2, 4, 8]) + # Attainable accuracy is floored by the condition number, not by the iteration count. + @parametrize("condition_number, tolerance", [(1e2, 1e-5), (1e4, 1e-4), (1e6, 1e-2)]) + def test_newton_schulz_root_inverse_matches_eigen( + self, condition_number: float, tolerance: float, root: int + ) -> None: + A = NewtonSchulzRootInverseTest._spd_matrix( + n=64, condition_number=condition_number + ) + self.assertLessEqual( + NewtonSchulzRootInverseTest._relative_error( + # relative_epsilon is disabled so this measures the accuracy of the iteration itself + # rather than the extra regularization it applies by default. + matrix_inverse_root( + A=A, + root=Fraction(root), + root_inv_config=NewtonSchulzRootInvConfig(relative_epsilon=0.0), + ), + matrix_inverse_root( + A=A, root=Fraction(root), root_inv_config=EigenConfig() + ), + ), + tolerance, + ) + + def test_newton_schulz_root_inverse_coefficient_schedule(self) -> None: + """A per-iteration coefficient schedule, the form Polar Express supplies, is accepted.""" + A = NewtonSchulzRootInverseTest._spd_matrix(n=32, condition_number=1e4) + X = matrix_inverse_root( + A=A, + root=Fraction(4), + root_inv_config=NewtonSchulzRootInvConfig( + coefficients=[[3.0, -16.0 / 5.0, 6.0 / 5.0]] * 12 + ), + ) + self.assertTrue(torch.isfinite(X).all()) + + def test_newton_schulz_root_inverse_applies_epsilon(self) -> None: + # Singular matrix: without the epsilon ridge the inverse root does not exist. + A = torch.tensor([[1.0, 0.0], [0.0, 0.0]]) + epsilon = 1e-2 + torch.testing.assert_close( + matrix_inverse_root( + A=A, + root=Fraction(2), + root_inv_config=NewtonSchulzRootInvConfig(), + epsilon=epsilon, + ), + matrix_inverse_root( + A=A, + root=Fraction(2), + root_inv_config=EigenConfig(), + epsilon=epsilon, + ), + atol=1e-3, + rtol=1e-3, + ) + + @parametrize("root", [1, 3, 6]) + def test_newton_schulz_root_inverse_unsupported_root(self, root: int) -> None: + self.assertRaisesRegex( + ValueError, + re.escape( + f"root={root} must be a power of two to use Newton-Schulz iteration!" + ), + matrix_inverse_root, + A=torch.eye(2), + root=Fraction(root), + root_inv_config=NewtonSchulzRootInvConfig(), + ) + + def test_newton_schulz_root_inverse_non_integer_root(self) -> None: + self.assertRaisesRegex( + ValueError, + re.escape( + "root.denominator=3 must be equal to 1 to use Newton-Schulz iteration!" + ), + matrix_inverse_root, + A=torch.tensor([[1.0, 0.0], [0.0, 4.0]]), + root=Fraction(2, 3), + root_inv_config=NewtonSchulzRootInvConfig(), + ) + + def test_newton_schulz_root_inverse_empty_coefficients(self) -> None: + self.assertRaisesRegex( + ValueError, + re.escape("coefficients must be non-empty."), + NewtonSchulzRootInvConfig, + coefficients=[], + ) + + def test_newton_schulz_root_inverse_warns_on_coefficients_not_summing_to_one( + self, + ) -> None: + with self.assertLogs(level="WARNING") as cm: + NewtonSchulzRootInvConfig(coefficients=[[3.4445, -4.7750, 2.0315]]) + self.assertIn("do not sum to 1", "".join(r.msg for r in cm.records)) + + @parametrize( + "coefficients", + [ + [[1.875, -1.25]], + [[1.875, -1.25, 0.375, 0.0]], + [[1.875, -1.25, 0.375], [1.0, 0.0]], + ], + ) + def test_newton_schulz_root_inverse_coefficients_wrong_arity( + self, coefficients: list[list[float]] + ) -> None: + self.assertRaisesRegex( + ValueError, + re.escape("must contain exactly three coefficients"), + NewtonSchulzRootInvConfig, + coefficients=coefficients, + ) + + @parametrize("coefficient", [float("nan"), float("inf"), float("-inf")]) + def test_newton_schulz_root_inverse_non_finite_coefficients( + self, coefficient: float + ) -> None: + self.assertRaisesRegex( + ValueError, + re.escape("must be finite"), + NewtonSchulzRootInvConfig, + coefficients=[[1.875, -1.25, coefficient]], + ) + + @parametrize("coefficient", ["0.375", None, True]) + def test_newton_schulz_root_inverse_non_numeric_coefficients( + self, coefficient: object + ) -> None: + self.assertRaisesRegex( + ValueError, + re.escape("must contain real numbers"), + NewtonSchulzRootInvConfig, + coefficients=[[1.875, -1.25, coefficient]], + ) + + def test_newton_schulz_root_inverse_accepts_integer_coefficients(self) -> None: + """Integers are real numbers; the arity/finiteness checks must not reject them.""" + self.assertTrue( + torch.isfinite( + matrix_inverse_root( + A=torch.eye(4), + root=Fraction(2), + root_inv_config=NewtonSchulzRootInvConfig( + coefficients=[[3, -3, 1]] + ), + ) + ).all() + ) + + @parametrize("dtype", [torch.float32, torch.float64]) + def test_newton_schulz_root_inverse_accepts_supported_dtype( + self, dtype: torch.dtype + ) -> None: + X = matrix_inverse_root( + A=NewtonSchulzRootInverseTest._spd_matrix(n=16, condition_number=1e2).to( + dtype=dtype + ), + root=Fraction(2), + root_inv_config=NewtonSchulzRootInvConfig(), + ) + self.assertIs(X.dtype, dtype) + self.assertTrue(torch.isfinite(X).all()) + + class CoupledHigherOrderRootInverseTest(unittest.TestCase): def test_root_with_big_numerator_denominator(self) -> None: A = torch.tensor([[1.0, 0.0], [0.0, 4.0]]) diff --git a/distributed_shampoo/shampoo_types.py b/distributed_shampoo/shampoo_types.py index 69f8864..f418096 100644 --- a/distributed_shampoo/shampoo_types.py +++ b/distributed_shampoo/shampoo_types.py @@ -12,6 +12,7 @@ from collections.abc import Callable from dataclasses import dataclass, field, make_dataclass from inspect import signature +from typing import Any import torch from distributed_shampoo.preconditioner.matrix_functions_types import ( @@ -1235,7 +1236,11 @@ class HybridShardDistributedConfig(FullyShardDistributedConfig, DDPDistributedCo device_mesh: DeviceMesh -_ShampooPT2CompileConfigImpl: type[object] = make_dataclass( +# Any, not type[object]: the fields are synthesized at import time from +# torch.compile's signature, so a checker cannot resolve them. Any keeps the +# subclass's attribute access permissive; type[object] instead made the class +# statically unusable as a base. +_ShampooPT2CompileConfigImpl: Any = make_dataclass( "_ShampooPT2CompileConfigImpl", [ (name, param.annotation, param.default) @@ -1246,9 +1251,7 @@ class HybridShardDistributedConfig(FullyShardDistributedConfig, DDPDistributedCo ) -class ShampooPT2CompileConfig( - _ShampooPT2CompileConfigImpl # type: ignore -): +class ShampooPT2CompileConfig(_ShampooPT2CompileConfigImpl): """Configuration for Shampoo PT2 compilation. Enables Shampoo pytorch compilation with configure to speed up model training. diff --git a/distributed_shampoo/tests/distributed_shampoo_test.py b/distributed_shampoo/tests/distributed_shampoo_test.py index 2263b9c..aa6fc46 100644 --- a/distributed_shampoo/tests/distributed_shampoo_test.py +++ b/distributed_shampoo/tests/distributed_shampoo_test.py @@ -21,6 +21,7 @@ from distributed_shampoo.preconditioner.matrix_functions_types import ( DefaultNewtonSchulzOrthogonalizationConfig, EigenConfig, + NewtonSchulzRootInvConfig, OrthogonalizationConfig, PseudoInverseConfig, ) @@ -71,6 +72,50 @@ def _pack_if_enabled( ) +@instantiate_parametrized_tests +class DistributedShampooNewtonSchulzTest(unittest.TestCase): + """Optimizer-level guards for the Newton-Schulz inverse root. + + The numerics live in NewtonSchulzRootInverseTest (preconditioner/tests/matrix_functions_test.py); + what can only be checked here is that an unsupported root is rejected while constructing the + optimizer, since the block orders that determine the root come from DistributedShampoo's + parameter blocking.""" + + @staticmethod + def _optim_factory( + parameters: Any, + preconditioner_config: PreconditionerConfig, + ) -> torch.optim.Optimizer: + return DistributedShampoo( + parameters, + lr=0.01, + betas=(0.9, 0.999), + epsilon=1e-8, + max_preconditioner_dim=5, + precondition_frequency=1, + start_preconditioning_step=1, + preconditioner_config=preconditioner_config, + ) + + def test_unsupported_root_fails_fast(self) -> None: + """An order-3 block asks for root 6, which Newton-Schulz cannot compute. This must raise at + optimizer construction rather than be swallowed as a per-factor-matrix warning during + training that silently reuses a stale preconditioner.""" + model = nn.ParameterList([nn.Parameter(torch.randn(4, 4, 4))]) + self.assertRaisesRegex( + ValueError, + re.escape( + "NewtonSchulzRootInvConfig only supports inverse roots that are powers of two, but " + "unsupported_roots=[6.0] were requested." + ), + DistributedShampooNewtonSchulzTest._optim_factory, + model.parameters(), + preconditioner_config=RootInvShampooPreconditionerConfig( + amortized_computation_config=NewtonSchulzRootInvConfig() + ), + ) + + @instantiate_parametrized_tests class DistributedShampooInitTest(unittest.TestCase): def setUp(self) -> None: diff --git a/distributed_shampoo/tests/shampoo_types_test.py b/distributed_shampoo/tests/shampoo_types_test.py index c255e87..e204d69 100644 --- a/distributed_shampoo/tests/shampoo_types_test.py +++ b/distributed_shampoo/tests/shampoo_types_test.py @@ -9,6 +9,8 @@ import re import unittest +from dataclasses import asdict, fields +from inspect import signature from typing import Any from unittest.mock import MagicMock @@ -31,6 +33,7 @@ HybridShardDistributedConfig, IterateAveragingConfig, RMSpropPreconditionerConfig, + ShampooPT2CompileConfig, SignDescentPreconditionerConfig, ) from distributed_shampoo.utils.commons import get_all_non_abstract_subclasses @@ -421,3 +424,24 @@ def test_illegal_num_sub_groups(self, num_sub_groups: int) -> None: device_mesh=MagicMock(), num_sub_groups=num_sub_groups, ) + + +class ShampooPT2CompileConfigTest(unittest.TestCase): + # The fields are synthesized at import time from torch.compile's signature, so no + # static checker can see them and nothing else in the suite asserts they exist. + def test_fields_match_torch_compile_signature(self) -> None: + self.assertEqual( + {field.name for field in fields(ShampooPT2CompileConfig())}, + {name for name in signature(torch.compile).parameters if name != "model"}, + ) + + def test_asdict_binds_to_torch_compile(self) -> None: + config = ShampooPT2CompileConfig(backend="eager", fullgraph=True) + kwargs = asdict(config) + self.assertEqual(kwargs["backend"], "eager") + self.assertTrue(kwargs["fullgraph"]) + # Mirrors how distributed_shampoo.py splats the config into torch.compile. + signature(torch.compile).bind(torch.nn.Identity(), **kwargs) + + def test_unknown_keyword_rejected(self) -> None: + self.assertRaises(TypeError, ShampooPT2CompileConfig, not_a_torch_compile_arg=1) diff --git a/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py b/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py index 994eb55..45b27ad 100644 --- a/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py +++ b/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py @@ -419,7 +419,7 @@ def test_param_recv_info_completeness(self) -> None: # Every param should have a valid recv info entry (offset >= 0) for param_idx in range(num_params): offset, chunk_size = ctx._param_recv_info[param_idx] - self.assertGreaterEqual( # type: ignore + self.assertGreaterEqual( offset, 0, f"Param {param_idx} has invalid recv offset {offset}" ) self.assertGreaterEqual( @@ -545,7 +545,7 @@ def test_gather_gradients_matches_full_tensor(self, num_params: int) -> None: # Verify correctness for assigned params for i in range(num_params): if i % self.world_size == rank: - self.assertIsNotNone( # type: ignore + self.assertIsNotNone( gathered_grads[i], f"Assigned param {i} should have a gathered gradient", ) @@ -556,7 +556,7 @@ def test_gather_gradients_matches_full_tensor(self, num_params: int) -> None: ) else: # Unassigned params should be None - self.assertIsNone( # type: ignore + self.assertIsNone( gathered_grads[i], f"Unassigned param {i} should be None on rank {rank}", ) @@ -589,9 +589,7 @@ def test_gather_gradients_with_none_grads(self) -> None: # All should be None since no gradients were set for i, grad in enumerate(gathered_grads): - self.assertIsNone( # type: ignore - grad, f"Param {i} has no grad, should be None" - ) + self.assertIsNone(grad, f"Param {i} has no grad, should be None") @with_comms @skip_if_lt_x_gpu(4) @@ -639,14 +637,14 @@ def test_gather_gradients_with_partial_none_grads(self) -> None: for i in range(len(shapes)): if i % self.world_size == rank: if expected_full_grads[i] is not None: - self.assertIsNotNone(gathered_grads[i]) # type: ignore + self.assertIsNotNone(gathered_grads[i]) torch.testing.assert_close( gathered_grads[i], expected_full_grads[i] ) else: - self.assertIsNone(gathered_grads[i]) # type: ignore + self.assertIsNone(gathered_grads[i]) else: - self.assertIsNone(gathered_grads[i]) # type: ignore + self.assertIsNone(gathered_grads[i]) @with_comms @skip_if_lt_x_gpu(4) @@ -673,7 +671,7 @@ def test_gather_gradients_preserves_shape(self) -> None: for i in range(len(shapes)): if i % self.world_size == rank: - self.assertIsNotNone(gathered_grads[i]) # type: ignore + self.assertIsNotNone(gathered_grads[i]) self.assertEqual( gathered_grads[i].shape, # type: ignore torch.Size(shapes[i]), @@ -722,11 +720,11 @@ def test_gather_gradients_multiple_calls(self) -> None: # Verify second call produces correct results (different from first) for i in range(len(shapes)): if i % self.world_size == rank: - self.assertIsNotNone(second_grads[i]) # type: ignore + self.assertIsNotNone(second_grads[i]) torch.testing.assert_close(second_grads[i], expected_second_grads[i]) # Verify second call differs from first # (2x gradient vs 1x gradient for sum) - self.assertFalse( # type: ignore + self.assertFalse( torch.equal(first_grads[i], second_grads[i]), # type: ignore f"Second gather should differ from first for param {i}", ) @@ -823,7 +821,7 @@ def test_gather_params_matches_full_tensor(self, num_params: int) -> None: # Verify correctness for assigned params for i in range(num_params): if i % self.world_size == rank: - self.assertIsNotNone( # type: ignore + self.assertIsNotNone( gathered_params[i], f"Assigned param {i} should have a gathered value", ) @@ -833,7 +831,7 @@ def test_gather_params_matches_full_tensor(self, num_params: int) -> None: msg=f"Gathered param {i} does not match full_tensor()", ) else: - self.assertIsNone( # type: ignore + self.assertIsNone( gathered_params[i], f"Unassigned param {i} should be None on rank {rank}", ) diff --git a/distributed_shampoo/utils/optimizer_modules.py b/distributed_shampoo/utils/optimizer_modules.py index 5de1ee9..40d33b9 100644 --- a/distributed_shampoo/utils/optimizer_modules.py +++ b/distributed_shampoo/utils/optimizer_modules.py @@ -87,6 +87,7 @@ def save_to_state_dict( for key, value in states: if isinstance(value, torch.Tensor): + # pyrefly: ignore [bad-argument-type] destination[key] = value if keep_vars else value.detach() elif isinstance(value, OptimizerModule): destination[key] = {} @@ -220,7 +221,6 @@ def load_from_new_state_to_old_state( old_state = type(old_state)( ( load_from_new_state_to_old_state( - # pyrefly: ignore [bad-argument-type] old_state=old_value, # pyrefly: ignore [bad-index] new_state=new_state[i],