[preconditioner] thread-safe TF32 context manager, generalize rank_deficient_stability_config to MatrixFunctionConfig, refactor _precondition_grad

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(

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:
@@ -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: "")

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.

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)