From bb4c0e8e743dd98121f9c31573e56c55e94332d0 Mon Sep 17 00:00:00 2001 From: Iris Zhang Date: Wed, 19 Aug 2026 20:50:18 -0700 Subject: [PATCH 1/2] Sync 3 diffs to facebookresearch/optimizers 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 was dropped. The `cast(torch.Tensor, param)` mypy fix (D113741320) already shipped to the mirror in an earlier sync and reappeared 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. 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. Also advances the non-exported last-sync marker at `scripts/.last_oss_sync_rev` to `084ddd3b2798` (D116658375), the newest committed dev-path commit. This covers every commit whose change is in this mirror, so the next sync will not re-list any of them. Differential Revision: D116726514 --- CHANGELOG.md | 8 +- distributed_shampoo/__init__.py | 2 + distributed_shampoo/distributed_shampoo.py | 3 +- distributed_shampoo/examples/utils.py | 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 | 115 ++++++++++ .../tests/shampoo_types_test.py | 24 +++ .../utils/gpu_tests/shampoo_utils_test.py | 26 ++- .../utils/optimizer_modules.py | 2 +- 13 files changed, 610 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 31e8e923..3e8b02e4 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 (blocks of order 1, 2, or 4); 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 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/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..4bd9c541 100644 --- a/distributed_shampoo/tests/distributed_shampoo_test.py +++ b/distributed_shampoo/tests/distributed_shampoo_test.py @@ -14,6 +14,7 @@ import unittest from collections.abc import Callable from dataclasses import dataclass, field, replace +from functools import partial from typing import Any, cast import torch @@ -21,6 +22,7 @@ from distributed_shampoo.preconditioner.matrix_functions_types import ( DefaultNewtonSchulzOrthogonalizationConfig, EigenConfig, + NewtonSchulzRootInvConfig, OrthogonalizationConfig, PseudoInverseConfig, ) @@ -51,6 +53,11 @@ TRAIN_MODE, WeightDecayType, ) +from distributed_shampoo.tests.shampoo_test_utils import ( + compare_two_optimizers_on_weight_and_loss, + construct_training_problem, + train_model, +) from distributed_shampoo.utils.shampoo_utils import pack_upper_triangular from torch import nn, Tensor from torch.testing._internal.common_utils import ( @@ -71,6 +78,114 @@ def _pack_if_enabled( ) +@instantiate_parametrized_tests +class DistributedShampooNewtonSchulzTest(unittest.TestCase): + """End-to-end training with the Newton-Schulz inverse root instead of an eigendecomposition.""" + + @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 _assert_trains(self, preconditioner_config: PreconditionerConfig) -> None: + model, loss, data, target, _ = train_model( + optim_factory=partial( + DistributedShampooNewtonSchulzTest._optim_factory, + preconditioner_config=preconditioner_config, + ), + model_factory=partial( + construct_training_problem, + model_linear_layers_dims=(10, 5, 3), + model_dead_layers_dims=None, + fill=0.1, + ), + num_steps=20, + ) + for name, parameter in model.named_parameters(): + self.assertTrue( + torch.isfinite(parameter).all(), + msg=f"Non-finite values in {name} after training.", + ) + # The problem targets zero, so a working preconditioner must drive the loss below its + # value at initialization. + self.assertLess(loss(model(data), target).item(), 0.1) + + def test_training_with_newton_schulz(self) -> None: + self._assert_trains( + RootInvShampooPreconditionerConfig( + amortized_computation_config=NewtonSchulzRootInvConfig() + ) + ) + + def test_training_with_coefficient_schedule(self) -> None: + """A per-iteration coefficient schedule, the form Polar Express supplies, trains end to end.""" + self._assert_trains( + RootInvShampooPreconditionerConfig( + amortized_computation_config=NewtonSchulzRootInvConfig( + coefficients=[[3.4445, -4.7750, 2.0315]] * 6 + + [[3.0, -16.0 / 5.0, 6.0 / 5.0]] * 14 + ) + ) + ) + + def test_matches_eigendecomposition_on_full_rank_factor_matrices(self) -> None: + """Newton-Schulz agrees with the eigendecomposition once the factor matrices are full rank. + + It does NOT agree while they are rank deficient, because relative_epsilon regularizes more + aggressively there than epsilon does on the eigendecomposition path. + """ + compare_two_optimizers_on_weight_and_loss( + control_optim_factory=partial( + DistributedShampooNewtonSchulzTest._optim_factory, + preconditioner_config=DefaultShampooConfig, + ), + experimental_optim_factory=partial( + DistributedShampooNewtonSchulzTest._optim_factory, + preconditioner_config=RootInvShampooPreconditionerConfig( + # Match the eigendecomposition path's regularization so the two are comparable. + amortized_computation_config=NewtonSchulzRootInvConfig( + relative_epsilon=0.0 + ) + ), + ), + model_linear_layers_dims=(10, 10), + model_dead_layers_dims=None, + fill=0.1, + total_steps=5, + rtol=1e-2, + atol=1e-3, + ) + + 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], From b5304b27f73f8454406b1a7fe42c17e05d9f7b23 Mon Sep 17 00:00:00 2001 From: Kang An Date: Fri, 21 Aug 2026 14:35:11 -0700 Subject: [PATCH 2/2] Sync 1 diff to facebookresearch/optimizers 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] Add per-head block splitting for attention projections, via two new `DDPDistributedConfig` fields: `head_split_size` and `head_split_param_names`. When `head_split_size > 0`, a matching 2D parameter is split along dim 0 into `head_split_size`-row blocks instead of taking the default merge-and-block path, so a matrix preconditioner (Shampoo or spectral descent) sees one block per attention head rather than the whole fused projection. A parameter is split iff its name contains any substring in `head_split_param_names`; parameters in the same group that do not match keep the default blocking. Opt-in -- `head_split_size` defaults to 0, so the default path is unchanged. Internal revisions: D115233938 **Most of this diff is machine output; one file is not.** The mirrored `.py` files under `distributed_shampoo/` are regenerated by the script. `CHANGELOG.md` was hand-edited: the generator emits ` -- `, and that bullet was replaced with the release note above. No other hand edits. The mirror content is exactly the three source files from the internal revision plus the changelog entry -- `distributor/shampoo_distributor.py`, `distributor/shampoo_ddp_distributor.py`, `shampoo_types.py`, and `CHANGELOG.md` (141 insertions, 15 deletions total). Note on the marker: `scripts/.last_oss_sync_rev` was advanced by the script to the internal revision's commit hash, which is still a draft. Its hash will be rewritten on every amend and again at land, so the next sync will not find it and will need an explicit `--since`. Setting it back to the newest landed commit instead would trade that hard failure for a soft one (the next sync re-lists this revision and generates a duplicate bullet). Flagging rather than hand-editing, since the marker is script-owned. Differential Revision: D116989675 --- CHANGELOG.md | 4 + .../distributor/shampoo_ddp_distributor.py | 42 ++++++++ .../distributor/shampoo_distributor.py | 96 ++++++++++++++++--- distributed_shampoo/shampoo_types.py | 14 +++ 4 files changed, 141 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e8b02e4..ae5927d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ Release notes for the OSS mirror. Seeded by `scripts/oss_sync.sh` from the internal source tree and edited for release. +## 2026-08-21 + +- [shampoo] Add per-head block splitting for attention projections, via two new `DDPDistributedConfig` fields: `head_split_size` and `head_split_param_names`. When `head_split_size > 0`, a matching 2D parameter is split along dim 0 into `head_split_size`-row blocks instead of taking the default merge-and-block path, so a matrix preconditioner (Shampoo or spectral descent) sees one block per attention head rather than the whole fused projection. A parameter is split iff its name contains any substring in `head_split_param_names`; parameters in the same group that do not match keep the default blocking. Opt-in — `head_split_size` defaults to 0, so the default path is unchanged. Requires per-parameter names under `param_names` in the param group, and raises if they are missing or their count does not match `params`; a negative `head_split_size` raises at config construction, and at block time a non-2D parameter, or a size that does not divide the parameter's first dimension, raises. If `head_split_size` is set but no name matches, a warning is logged and the split is a no-op. Implemented for the DDP distributor. + ## 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 (blocks of order 1, 2, or 4); 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`). diff --git a/distributed_shampoo/distributor/shampoo_ddp_distributor.py b/distributed_shampoo/distributor/shampoo_ddp_distributor.py index df386423..5acab4a9 100644 --- a/distributed_shampoo/distributor/shampoo_ddp_distributor.py +++ b/distributed_shampoo/distributor/shampoo_ddp_distributor.py @@ -24,6 +24,7 @@ from distributed_shampoo.shampoo_types import ( DDPDistributedConfig, DISTRIBUTED_CONFIG, + PARAM_NAMES, PARAMS, ShampooRuntimeConfig, ) @@ -220,6 +221,47 @@ def __init__( group_rank=group_rank, ) + def _resolve_head_split_sizes(self) -> tuple[int | None, ...] | None: + """Resolve per-parameter head-split sizes from the DDP config. + + Returns None (head-split disabled) unless ``head_split_size > 0``. When + enabled, a parameter is head-split iff its name contains any substring in + ``head_split_param_names``; ``head_split_size`` (head_dim) is used for + those and None for the rest. Requires per-parameter names under + PARAM_NAMES in the param group. Works with any matrix preconditioner + (e.g. Shampoo or spectral descent / Muon). + """ + distributed_config: DDPDistributedConfig = self._param_group[DISTRIBUTED_CONFIG] + head_split_size = distributed_config.head_split_size + if not head_split_size: + return None + if PARAM_NAMES not in self._param_group: + raise ValueError( + f"head_split_size requires per-parameter names under " + f"'{PARAM_NAMES}' in the param group." + ) + param_names = self._param_group[PARAM_NAMES] + num_params = len(self._param_group[PARAMS]) + if len(param_names) != num_params: + raise ValueError( + f"len({PARAM_NAMES})={len(param_names)} must equal " + f"len(params)={num_params}." + ) + head_split_sizes = tuple( + head_split_size + if any( + pattern in name for pattern in distributed_config.head_split_param_names + ) + else None + for name in param_names + ) + if not any(size is not None for size in head_split_sizes): + logger.warning( + "head_split_size is set but no parameter name matched " + "head_split_param_names; head-split is a no-op." + ) + return head_split_sizes + @overload @torch.no_grad() def _get_params_or_grads( diff --git a/distributed_shampoo/distributor/shampoo_distributor.py b/distributed_shampoo/distributor/shampoo_distributor.py index 645901db..4bcfd8ba 100644 --- a/distributed_shampoo/distributor/shampoo_distributor.py +++ b/distributed_shampoo/distributor/shampoo_distributor.py @@ -7,8 +7,9 @@ """ +import logging from abc import ABC, abstractmethod -from collections.abc import Iterable +from collections.abc import Callable, Iterable from functools import partial from operator import attrgetter from typing import Any, Literal, overload @@ -29,6 +30,8 @@ ) from torch import Tensor +logger: logging.Logger = logging.getLogger(__name__) + ###### DISTRIBUTOR CLASSES ###### class DistributorInterface(ABC): @@ -51,6 +54,13 @@ def __init__( self._runtime_config: ShampooRuntimeConfig = ( runtime_config if runtime_config is not None else ShampooRuntimeConfig() ) + # Per-parameter head-split sizes (None disables head-split for every + # param). Resolved via virtual dispatch so only distributors that + # support head-split (currently DDP) return a non-None spec. Must be set + # before _merge_and_block_parameters() so blocking can consult it. + self._head_split_sizes: tuple[int | None, ...] | None = ( + self._resolve_head_split_sizes() + ) # Merge and block parameters creates self._global_blocked_params and self._global_num_blocks_per_param # Global blocked params are all the blocked parameters after merging and blocking. # Global num blocks per param stores the number of blocks for each global parameter. @@ -75,6 +85,21 @@ def __init__( # Local block info list contains information about each block masked by the distributor selector. self._local_block_info_list: tuple[BlockInfo, ...] + def _resolve_head_split_sizes(self) -> tuple[int | None, ...] | None: + """Return per-parameter head-split sizes, or None to disable head-split. + + Base implementation disables head-split. Distributors that support it + (currently only ``DDPDistributor``) override this. A warning is emitted + if head-split was requested on an unsupported distributor so the config + is not silently ignored. + """ + if getattr(self._param_group[DISTRIBUTED_CONFIG], "head_split_size", 0): + logger.warning( + "head_split_size is set but ignored: head-split is only " + f"supported by DDPDistributor, not {type(self).__name__}." + ) + return None + @abstractmethod @torch.no_grad() def update_params( @@ -170,6 +195,43 @@ def _get_params_or_grads(self, get_grad: bool = False) -> Iterable[Tensor | None else self._param_group[PARAMS] ) + def _blocks_within_tensor( + self, + tensor: Tensor, + param_index: int, + merge_dims: Callable[..., tuple[int, ...]], + ) -> tuple[Tensor, ...]: + """Split one param/grad tensor into preconditioner blocks. + + Default path: merge small dims, then block by ``max_preconditioner_dim``. + When head-split is enabled for this param, instead split dim 0 into + ``head_dim``-sized 2D blocks (one per attention head) and skip merging / + blocking so the model (input) dimension is preserved. Both the parameter + and gradient paths call this so their block decomposition stays + identical. + """ + head_dim = ( + None + if self._head_split_sizes is None + else self._head_split_sizes[param_index] + ) + if head_dim is not None: + if tensor.dim() != 2: + raise ValueError( + "head-split requires a 2D parameter, got shape " + f"{tuple(tensor.shape)} at param index {param_index}." + ) + if tensor.shape[0] % head_dim != 0: + raise ValueError( + f"head_split_size={head_dim} must divide out_features=" + f"{tensor.shape[0]} at param index {param_index}." + ) + return torch.split(tensor, head_dim, dim=0) + return multi_dim_split( + tensor.view(merge_dims(tensor_shape=tensor.size())), + self._param_group[MAX_PRECONDITIONER_DIM], + ) + @torch.no_grad() def _merge_and_block_with_params( self, params: Iterable[Tensor] @@ -199,11 +261,10 @@ def _merge_and_block_with_params( ].target_parameter_dimensionality, ) - for param in params: - # Obtain blocks for each parameter after merging. - blocks_within_param = multi_dim_split( - param.view(merge_dims(tensor_shape=param.size())), - self._param_group[MAX_PRECONDITIONER_DIM], + for param_index, param in enumerate(params): + # Obtain blocks for each parameter after merging (or per-head split). + blocks_within_param = self._blocks_within_tensor( + param, param_index, merge_dims ) # Generate and extend blocked parameters list. @@ -256,11 +317,17 @@ def _merge_and_block_gradients( ].target_parameter_dimensionality, ) - for grad, num_blocks, (block_index, next_block_index) in zip( - self._get_params_or_grads(get_grad=True), - self._global_num_blocks_per_param, - generate_pairwise_indices(self._global_num_blocks_per_param), - strict=True, + for param_index, ( + grad, + num_blocks, + (block_index, next_block_index), + ) in enumerate( + zip( + self._get_params_or_grads(get_grad=True), + self._global_num_blocks_per_param, + generate_pairwise_indices(self._global_num_blocks_per_param), + strict=True, + ) ): param_distributor_selector = self._distributor_selector[ block_index:next_block_index @@ -283,10 +350,9 @@ def _merge_and_block_gradients( f"Encountered gradient containing NaN/Inf in parameter with shape {attrgetter('shape')(grad)}. Check your model for numerical instability or consider gradient clipping." ) - # Obtain blocks for each gradient after merging. - blocks_within_grad = multi_dim_split( - grad.view(merge_dims(tensor_shape=grad.size())), - self._param_group[MAX_PRECONDITIONER_DIM], + # Obtain blocks for each gradient after merging (or per-head split). + blocks_within_grad = self._blocks_within_tensor( + grad, param_index, merge_dims ) # Generate block-to-parameter metadata and extend blocked parameters list. local_masked_blocked_grads.extend( diff --git a/distributed_shampoo/shampoo_types.py b/distributed_shampoo/shampoo_types.py index f4180969..caaa37d6 100644 --- a/distributed_shampoo/shampoo_types.py +++ b/distributed_shampoo/shampoo_types.py @@ -51,6 +51,7 @@ LR = "lr" MAX_PRECONDITIONER_DIM = "max_preconditioner_dim" PARAMS = "params" # While this is stored in groups by default, we do not checkpoint this quantity. +PARAM_NAMES = "param_names" PEAK_LR = "peak_lr" PRECONDITION_FREQUENCY = "precondition_frequency" PRECONDITIONER_CONFIG = "preconditioner_config" @@ -1086,6 +1087,12 @@ class DDPDistributedConfig(DistributedConfig): communication_dtype: torch.dtype = torch.float32 num_trainers_per_group: int = -1 communicate_params: bool = False + # Per-head splitting of attention QKV weights for Muon / spectral descent. + # head_split_size = head_dim (0 disables). head_split_param_names lists the + # FQN substrings (e.g. "qkv_linear.") of the projections to split along + # dim 0 into one 2D [head_dim, in_features] block per attention head. + head_split_size: int = 0 + head_split_param_names: list[str] = field(default_factory=list) @staticmethod def _get_default_load_balancing_config() -> LoadBalancingConfig: @@ -1095,6 +1102,13 @@ def _get_default_load_balancing_config() -> LoadBalancingConfig: default_factory=_get_default_load_balancing_config ) + def __post_init__(self) -> None: + super().__post_init__() + if self.head_split_size < 0: + raise ValueError( + f"Invalid {self.head_split_size=}. Must be >= 0 (0 disables head-split)." + ) + @dataclass(kw_only=True) class FSDPDistributedConfig(DistributedConfig):