From ba05cb939f80cb9fcee7e9e769ff10d12252c33a Mon Sep 17 00:00:00 2001 From: Iris Zhang Date: Mon, 24 Aug 2026 16:08:44 -0700 Subject: [PATCH] Sync 3 diffs to facebookresearch/optimizers (#276) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Syncs the internal dev tree to the OSS mirror (`fbcode/ai_codesign/optimizers`) via `scripts/oss_sync.sh`. Changes included since the last sync marker: - [shampoo] Newton-Schulz inverse root (D115751122) - [shampoo][dev] Fix OSS mypy on `ShampooPT2CompileConfig`'s synthesized base class (D116658375, below this diff in the stack) - [pyre][automated] Remove unused type error suppressions - hpc (D116361970) **Most of this diff is machine output; two files are not.** The mirrored `.py` files under `distributed_shampoo/` and `gpa/` are regenerated by the script. `CHANGELOG.md` and `distributed_shampoo/__init__.py` were edited by hand — details below. Every bullet in the new `## 2026-08-19` entry was hand-written. The generator emits ` — `, which for these commits produced a link preamble, internal tooling narration, and a sentence truncated mid-clause. The generator produced four bullets; one is not in the current entry. The `cast(torch.Tensor, param)` mypy fix (D113741320) already shipped to the mirror in D113741175 and reappeared in this range only because that sync left the marker short. Verified against the diff rather than the title: the `distributed_shampoo_test.py` hunks here are all additions for the new Newton-Schulz coverage, with no `cast` lines. That commit turned out to be documented nowhere — it shipped publicly and no entry mentioned it. So instead of dropping it, this diff **backfills** a bullet into `## 2026-07-24`, the entry for the release that actually shipped it (`dd8c84ce4356`). That is the second `CHANGELOG.md` hand edit here and the only change outside the current entry. The `[pyre][automated]` sweep **does** get a bullet, despite the internal-looking title. It is not internal-only in the mirror: the transform stage rewrites internal `# pyre-ignore` into `# type: ignore`, so the sweep removes 13 `# type: ignore` comments plus a net 2 `# pyrefly: ignore` comments from the public tree. A reader diffing two releases would otherwise find unexplained changes. The bullet keeps the fact and drops the internal framing. `distributed_shampoo/__init__.py` is the one mirror `.py` file `oss_sync.sh` does not regenerate (it is manually maintained, by design). `NewtonSchulzRootInvConfig` was missing from it while the release note advertises the feature, so it is added to the imports and `__all__` alongside its sibling `RootInvConfig` subclasses. No sync stage catches this class of omission — mypy and all four CPU suites pass either way, because nothing in the mirror imports from the package root. The non-exported last-sync marker at `scripts/.last_oss_sync_rev` is set to `efefb9b4` (D115751122), the newest **landed** commit in the synced range. The script had advanced it to the dev fix stacked below (D116658375), which is the newest commit whose change is in this mirror and would be the right answer — except that commit is a draft, so its hash is rewritten by every amend and again at land. It was already off the tip's ancestry, which makes the next sync abort with `error: base rev ... was not found` rather than merely re-listing something. Falling back to the last landed commit trades a hard failure for a soft one: the next sync re-lists D116658375 and generates a duplicate bullet, which the oss-sync skill drops on the "already shipped in an earlier sync" rule. No post-land follow-up is required, and there is nothing to remember. (`--summary-only --since` is not an alternative — that mode copies no code, so `advance_marker` deliberately refuses to move the marker at all.) Reviewed By: anana10c Differential Revision: D116726514 --- CHANGELOG.md | 9 +- distributed_shampoo/__init__.py | 2 + distributed_shampoo/distributed_shampoo.py | 3 +- distributed_shampoo/examples/utils.py | 1 - distributed_shampoo/preconditioner/README.md | 1 + .../preconditioner/matrix_functions.py | 133 ++++++++++++ .../preconditioner/matrix_functions_types.py | 84 ++++++++ .../shampoo_preconditioner_list.py | 21 ++ .../tests/matrix_functions_test.py | 203 ++++++++++++++++++ distributed_shampoo/shampoo_types.py | 11 +- .../tests/distributed_shampoo_test.py | 45 ++++ .../tests/shampoo_types_test.py | 24 +++ .../utils/gpu_tests/shampoo_utils_test.py | 26 ++- .../utils/optimizer_modules.py | 2 +- 14 files changed, 542 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 31e8e923..4a1895a3 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 767cf6a6..2ec751df 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 0bbbc844..bdf8014c 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 a237c932..291a2949 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 de5dbf0f..bad331c1 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 48347e41..a67c8de9 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 02e7c63c..ade8cf30 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 d4a67329..db577175 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 0ccf5b7a..296e051c 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 69f88649..f4180969 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 2263b9cd..aa6fc46a 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 c255e871..e204d69e 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 994eb55b..45b27ade 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 5de1ee9e..40d33b9e 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],