From 5e649e90cdaf82ce85997b682a1192bd8a656e5f Mon Sep 17 00:00:00 2001 From: "Vishal S. Pandey" Date: Mon, 3 Aug 2026 13:21:31 +0530 Subject: [PATCH 1/2] preconditioner: thread-safe TF32, generalize rank_deficient_stability_config, and optimize _precondition_grad --- .../preconditioner/matrix_functions.py | 50 +++++++++--------- .../preconditioner/matrix_functions_types.py | 34 ++++++------- .../shampoo_preconditioner_list.py | 51 ++++++++----------- .../tests/matrix_functions_test.py | 21 ++++++++ 4 files changed, 86 insertions(+), 70 deletions(-) diff --git a/distributed_shampoo/preconditioner/matrix_functions.py b/distributed_shampoo/preconditioner/matrix_functions.py index 48347e41..e09ce512 100644 --- a/distributed_shampoo/preconditioner/matrix_functions.py +++ b/distributed_shampoo/preconditioner/matrix_functions.py @@ -11,8 +11,10 @@ import inspect import logging import math +import threading import time -from collections.abc import Callable +from collections.abc import Callable, Generator +from contextlib import contextmanager from dataclasses import fields from fractions import Fraction from functools import partial, wraps @@ -43,6 +45,27 @@ logger: logging.Logger = logging.getLogger(__name__) +_TF32_LOCK: threading.RLock = threading.RLock() + + +@contextmanager +def _scoped_tf32_setting(disable_tf32: bool) -> Generator[None, None, None]: + """Context manager to safely modify torch.backends.cuda.matmul.allow_tf32 in a thread-safe manner.""" + if not disable_tf32: + yield + return + + with _TF32_LOCK: + tf32_flag = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + logger.debug( + f"Using tf32 precision for fp32 matmul: {torch.backends.cuda.matmul.allow_tf32}" + ) + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = tf32_flag + @enum.unique class NewtonConvergenceFlag(enum.Enum): @@ -489,17 +512,7 @@ def matrix_inverse_root_higher_order( """ - # TODO(irisz): This save/modify/restore pattern is not thread-safe. - # Concurrent calls (e.g., from multi-threaded optimizer step) can race - # on this global flag. Revisit this for D97459682. - tf32_flag = torch.backends.cuda.matmul.allow_tf32 - if disable_tf32: - torch.backends.cuda.matmul.allow_tf32 = False - logger.debug( - f"Using tf32 precision for fp32 matmul: {torch.backends.cuda.matmul.allow_tf32}" - ) - - try: + with _scoped_tf32_setting(disable_tf32=disable_tf32): t_iter_begin = time.perf_counter() p = root.numerator q = root.denominator @@ -643,22 +656,12 @@ def matrix_inverse_root_higher_order( logger.debug(f"Error before powering: {true_error}") logger.debug(f"Termination Flag: {termination_flag}") - # If we have inf/nan in our answer also raise an arithmetic exception. - # Usually, this is due to the powering to q > 1 which can blow up entries. - # We have not seen this yet for q = 1 in Shampoo. if not torch.isfinite(X).all(): raise ArithmeticError( "NaN/Inf in matrix inverse root (after powering for fractions), raising an exception!" ) - 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. When - # disable_tf32=True, this restores the original value. - torch.backends.cuda.matmul.allow_tf32 = tf32_flag - - return X, M, termination_flag, iteration, true_error + return X, M, termination_flag, iteration, true_error match root_inv_config: case EigenConfig(): @@ -851,7 +854,6 @@ def qr_algorithm( return eigenvalues_estimate, eigenvectors_estimate - # TODO: reduce redundant code when rank_deficient_stability_config is generalized to all methods # check epsilon is 0 when using pseudo-inverse if ( isinstance( diff --git a/distributed_shampoo/preconditioner/matrix_functions_types.py b/distributed_shampoo/preconditioner/matrix_functions_types.py index 02e7c63c..35fedde4 100644 --- a/distributed_shampoo/preconditioner/matrix_functions_types.py +++ b/distributed_shampoo/preconditioner/matrix_functions_types.py @@ -32,7 +32,6 @@ class PerturbationConfig(RankDeficientStabilityConfig): that both options are mathematically equivalent, but not necessarily numerically equivalent. For eigenvalue-corrected Shampoo this will only affect the stability of the eigenbasis computation and epsilon will always also be added to the corrected eigenvalues. Recommended to be set to True for numerical stability. - TODO: When generalizing to all MatrixFunctionConfigs, this is only applicable to EigendecompositionConfig. (Default: True) """ @@ -50,11 +49,9 @@ class PseudoInverseConfig(RankDeficientStabilityConfig): Attributes: rank_atol: Absolute tolerance for filtering singular values. - TODO: When generalizing to all MatrixFunctionConfigs, this is only applicable to EigendecompositionConfig. (Default: 0.0) rank_rtol: Relative tolerance for filtering singular values. When None, takes value of max dim of the matrix times the epsilon of the dtype of the matrix. - TODO: When generalizing to all MatrixFunctionConfigs, this is only applicable to EigendecompositionConfig. (Default: 0.0) """ @@ -64,7 +61,23 @@ class PseudoInverseConfig(RankDeficientStabilityConfig): @dataclass(init=False) class MatrixFunctionConfig(AbstractDataclass): - """Base dataclass for matrix function configurations.""" + """Base dataclass for matrix function configurations. + + Note: When using custom rank_deficient_stability_config, avoid lambda functions as they may cause + pickling issues during serialization/deserialization. Use regular named functions + instead for better compatibility with distributed training and checkpointing. + + Attributes: + rank_deficient_stability_config (RankDeficientStabilityConfig): Configuration for handling/stabilizing rank-deficient matrices. (Default: DefaultPerturbationConfig) + """ + + @staticmethod + def _get_default_rank_deficient_stability_config() -> RankDeficientStabilityConfig: + return DefaultPerturbationConfig + + rank_deficient_stability_config: RankDeficientStabilityConfig = field( + default_factory=_get_default_rank_deficient_stability_config + ) @dataclass(init=False) @@ -84,25 +97,12 @@ class EigendecompositionConfig(MatrixFunctionConfig): Moreover, we have ||B||_F = ||Q^T A Q||_F = ||A||_F. Hence, the two relative errors are also equivalent: ||A - A'||_F / ||A||_F = ||B - diag(B)||_F / ||B||_F. - Note: When using custom rank_deficient_stability_config, avoid lambda functions as they may cause - pickling issues during serialization/deserialization. Use regular named functions - instead for better compatibility with distributed training and checkpointing. - Attributes: - rank_deficient_stability_config (RankDeficientStabilityConfig): Configuration for handling/stabilizing rank-deficient matrices. (Default: DefaultPerturbationConfig) - TODO: generalize this to MatrixFunctionConfig tolerance (float): The tolerance for the error of the eigendecomposition based on the norm of the off-diagonal elements of the eigenvalue estimate. (Default: 0.0) """ - @staticmethod - def _get_default_rank_deficient_stability_config() -> RankDeficientStabilityConfig: - return DefaultPerturbationConfig - - rank_deficient_stability_config: RankDeficientStabilityConfig = field( - default_factory=_get_default_rank_deficient_stability_config - ) tolerance: float = 0.0 def __post_init__(self) -> None: diff --git a/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py b/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py index d4a67329..02d73565 100644 --- a/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py +++ b/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py @@ -13,7 +13,7 @@ from collections.abc import Callable, Hashable, Mapping from dataclasses import asdict, dataclass, field, fields from fractions import Fraction -from functools import partial, reduce +from functools import partial from itertools import chain from operator import attrgetter from pathlib import Path @@ -1545,8 +1545,6 @@ def _precondition_grad( preconditioner_list: tuple[Tensor, ...], dims: tuple[list[int], list[int]] = ([0], [0]), ) -> Tensor: - # TODO: Need to refactor this function to be more efficient. Ideally eliminate those branches. - # Might consider einsum? assert sum(preconditioned_dims_selector) == len(preconditioner_list), ( f"The number of dimensions to precondition ({sum(preconditioned_dims_selector)}) must match the number of preconditioners ({len(preconditioner_list)})." ) @@ -1558,24 +1556,29 @@ def _precondition_grad( # Use the single dtype if preconditioners exist, otherwise use grad dtype target_dtype = next(iter(unique_dtypes), grad.dtype) - preconditioner_list_iter = iter(preconditioner_list) - - return reduce( - lambda grad, should_precondition: ( - torch.tensordot( - # Use the single target dtype for all operations - grad.to(dtype=target_dtype), - # Use the actual iterator for the operation - next(preconditioner_list_iter), + orig_dtype = grad.dtype + curr_grad = grad if orig_dtype == target_dtype else grad.to(dtype=target_dtype) + + p_idx = 0 + ndim = grad.ndim + permute_dims = (*range(1, ndim), 0) if ndim > 1 else () + + for should_precondition in preconditioned_dims_selector: + if should_precondition: + curr_grad = torch.tensordot( + curr_grad, + preconditioner_list[p_idx], dims=dims, ) - if should_precondition - # Perform a left rotation on grad if not preconditioned. - else grad.permute(*range(1, grad.ndim), 0) - ), - preconditioned_dims_selector, - grad, - ).to(dtype=grad.dtype) + p_idx += 1 + elif ndim > 1: + curr_grad = curr_grad.permute(permute_dims) + + return ( + curr_grad + if curr_grad.dtype == orig_dtype + else curr_grad.to(dtype=orig_dtype) + ) @overload @staticmethod @@ -1804,11 +1807,6 @@ def _compute_preconditioned_gradient( preconditioned_dims_selector: tuple[bool, ...], kronecker_factors: EigendecomposedShampooKroneckerFactorsUnwrapped, ) -> Tensor: - # TODO: remove assertion when rank_deficient_stability_config is generalized to MatrixFunctionConfig - assert isinstance( - self._preconditioner_config.amortized_computation_config, - EigendecompositionConfig, - ) rank_deficient_stability_config = self._preconditioner_config.amortized_computation_config.rank_deficient_stability_config return self._precondition_grad( @@ -1999,11 +1997,6 @@ def _compute_outer_product_list( preconditioned_dims_selector: tuple[bool, ...], kronecker_factors: EigendecomposedShampooKroneckerFactorsUnwrapped, ) -> tuple[Tensor, ...]: - # TODO: remove assertion when rank_deficient_stability_config is generalized to MatrixFunctionConfig - assert isinstance( - self._preconditioner_config.amortized_computation_config, - EigendecompositionConfig, - ) rank_deficient_stability_config = self._preconditioner_config.amortized_computation_config.rank_deficient_stability_config # Construct outer product list for updating Kronecker factors. diff --git a/distributed_shampoo/preconditioner/tests/matrix_functions_test.py b/distributed_shampoo/preconditioner/tests/matrix_functions_test.py index 0ccf5b7a..51b7f0ae 100644 --- a/distributed_shampoo/preconditioner/tests/matrix_functions_test.py +++ b/distributed_shampoo/preconditioner/tests/matrix_functions_test.py @@ -7,6 +7,7 @@ """ +import concurrent.futures import itertools import re import unittest @@ -23,6 +24,7 @@ _check_2d_tensor, _check_square_matrix, _matrix_perturbation, + _scoped_tf32_setting, matrix_eigendecomposition, matrix_inverse_root, matrix_inverse_root_from_eigendecomposition, @@ -960,3 +962,22 @@ class NotSupportedOrthogonalizationConfig(OrthogonalizationConfig): A=torch.tensor([[1.0, 0.0], [0.0, 4.0]]), orthogonalization_config=NotSupportedOrthogonalizationConfig(), ) + + def test_scoped_tf32_setting_thread_safety(self) -> None: + initial_flag = torch.backends.cuda.matmul.allow_tf32 + errors: list[Exception] = [] + + def worker(disable: bool) -> None: + try: + with _scoped_tf32_setting(disable_tf32=disable): + if disable: + self.assertFalse(torch.backends.cuda.matmul.allow_tf32) + except Exception as e: + errors.append(e) + + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: + futures = [executor.submit(worker, bool(i % 2)) for i in range(50)] + concurrent.futures.wait(futures) + + self.assertEqual(len(errors), 0) + self.assertEqual(torch.backends.cuda.matmul.allow_tf32, initial_flag) From 2ecc1051df2ce08b2fdcd582c0a9508de4f34e00 Mon Sep 17 00:00:00 2001 From: "Vishal S. Pandey" Date: Mon, 3 Aug 2026 13:34:34 +0530 Subject: [PATCH 2/2] attributes changed --- distributed_shampoo/preconditioner/matrix_functions_types.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/distributed_shampoo/preconditioner/matrix_functions_types.py b/distributed_shampoo/preconditioner/matrix_functions_types.py index 35fedde4..1e2800a5 100644 --- a/distributed_shampoo/preconditioner/matrix_functions_types.py +++ b/distributed_shampoo/preconditioner/matrix_functions_types.py @@ -129,7 +129,6 @@ class EighEigendecompositionConfig(EigendecompositionConfig): Hence, the two relative errors are also equivalent: ||A - A'||_F / ||A||_F = ||B - diag(B)||_F / ||B||_F. Attributes: - rank_deficient_stability_config (RankDeficientStabilityConfig): Configuration for handling/stabilizing rank-deficient matrices. (Default: DefaultPerturbationConfig) retry_double_precision (bool): Whether to retry eigendecomposition with higher (double) precision if lower precision fails due to CuSOLVER failure. (Default: True) eigendecomposition_offload_device (str): Device to offload eigendecomposition to. If value is empty string, we don't perform offloading. (Default: "") @@ -161,7 +160,6 @@ class QREigendecompositionConfig(EigendecompositionConfig): Hence, the two relative errors are also equivalent: ||A - A'||_F / ||A||_F = ||B - diag(B)||_F / ||B||_F. Attributes: - rank_deficient_stability_config (RankDeficientStabilityConfig): Configuration for handling/stabilizing rank-deficient matrices. (Default: DefaultPerturbationConfig) max_iterations (int): The maximum number of iterations to perform. (Default: 1) tolerance (float): The tolerance for determining convergence in terms of the norm of the off-diagonal elements of the eigenvalue estimate. (Default: 0.0) @@ -181,7 +179,6 @@ class EigenConfig(RootInvConfig, EighEigendecompositionConfig): """Configuration for matrix root inverse via an eigendecomposition. Attributes: - rank_deficient_stability_config (RankDeficientStabilityConfig): Configuration for handling/stabilizing rank-deficient matrices. (Default: DefaultPerturbationConfig) retry_double_precision (bool): Whether to retry eigendecomposition with higher (double) precision if lower precision fails due to CuSOLVER failure. (Default: True) eigendecomposition_offload_device (str): Device to offload eigendecomposition to. If value is empty string, we don't perform offloading. (Default: "")