Add control over tolerance for failed amortized computations

distributed_shampoo/shampoo_types.py

@@ -85,10 +85,18 @@ class PreconditionerConfig(AbstractDataclass):
 
     Args:
         amortized_computation_config (MatrixFunctionConfig): Configuration for the amortized computation, e.g., inverse-root or eigenvector computation.
+        num_tolerated_failed_amortized_computations (int): Number of failed amortized computations to tolerate before raising an error. (Default: 3)
 
     """
 
     amortized_computation_config: MatrixFunctionConfig  # type: ignore
+    num_tolerated_failed_amortized_computations: int = 3
+
+    def __post_init__(self) -> None:
+        if self.num_tolerated_failed_amortized_computations < 0:
+            raise ValueError(
+                f"Invalid num_tolerated_failed_amortized_computations value: {self.num_tolerated_failed_amortized_computations}. Must be >= 0."
+            )
 
 
 @dataclass(kw_only=True)

distributed_shampoo/tests/shampoo_types_test.py

@@ -9,24 +9,29 @@
 
 import re
 import unittest
-from typing import Type
+from abc import ABC, abstractmethod
+from typing import Generic, Type, TypeVar
 
 from distributed_shampoo.shampoo_types import (
     AdaGradGraftingConfig,
     AdamGraftingConfig,
+    EigenvalueCorrectedShampooPreconditionerConfig,
+    PreconditionerConfig,
     RMSpropGraftingConfig,
+    ShampooPreconditionerConfig,
 )
 
 
 class AdaGradGraftingConfigTest(unittest.TestCase):
     def test_illegal_epsilon(self) -> None:
         epsilon = 0.0
         grafting_config_type = self._get_grafting_config_type()
-        with self.subTest(
-            grafting_config_type=grafting_config_type
-        ), self.assertRaisesRegex(
-            ValueError,
-            re.escape(f"Invalid epsilon value: {epsilon}. Must be > 0.0."),
+        with (
+            self.subTest(grafting_config_type=grafting_config_type),
+            self.assertRaisesRegex(
+                ValueError,
+                re.escape(f"Invalid epsilon value: {epsilon}. Must be > 0.0."),
+            ),
         ):
             grafting_config_type(epsilon=epsilon)
 
@@ -46,12 +51,13 @@ def test_illegal_beta2(
     ) -> None:
         grafting_config_type = self._get_grafting_config_type()
         for beta2 in (-1.0, 0.0, 1.3):
-            with self.subTest(
-                grafting_config_type=grafting_config_type, beta2=beta2
-            ), self.assertRaisesRegex(
-                ValueError,
-                re.escape(
-                    f"Invalid grafting beta2 parameter: {beta2}. Must be in (0.0, 1.0]."
+            with (
+                self.subTest(grafting_config_type=grafting_config_type, beta2=beta2),
+                self.assertRaisesRegex(
+                    ValueError,
+                    re.escape(
+                        f"Invalid grafting beta2 parameter: {beta2}. Must be in (0.0, 1.0]."
+                    ),
                 ),
             ):
                 grafting_config_type(beta2=beta2)
@@ -67,3 +73,57 @@ def _get_grafting_config_type(
         self,
     ) -> Type[RMSpropGraftingConfig] | Type[AdamGraftingConfig]:
         return AdamGraftingConfig
+
+
+PreconditionerConfigType = TypeVar(
+    "PreconditionerConfigType", bound=Type[PreconditionerConfig]
+)
+
+
+class AbstractPreconditionerConfigTest:
+    class PreconditionerConfigTest(
+        ABC,
+        unittest.TestCase,
+        Generic[PreconditionerConfigType],
+    ):
+        def test_illegal_num_tolerated_failed_amortized_computations(self) -> None:
+            num_tolerated_failed_amortized_computations = -1
+            with (
+                self.assertRaisesRegex(
+                    ValueError,
+                    re.escape(
+                        f"Invalid num_tolerated_failed_amortized_computations value: "
+                        f"{num_tolerated_failed_amortized_computations}. Must be >= 0."
+                    ),
+                ),
+            ):
+                self._get_preconditioner_config_type()(
+                    num_tolerated_failed_amortized_computations=num_tolerated_failed_amortized_computations,
+                )
+
+        @abstractmethod
+        def _get_preconditioner_config_type(
+            self,
+        ) -> PreconditionerConfigType: ...
+
+
+class ShampooPreconditionerConfigTest(
+    AbstractPreconditionerConfigTest.PreconditionerConfigTest[
+        Type[ShampooPreconditionerConfig]
+    ]
+):
+    def _get_preconditioner_config_type(
+        self,
+    ) -> Type[ShampooPreconditionerConfig]:
+        return ShampooPreconditionerConfig
+
+
+class EigenvalueCorrectedShampooPreconditionerConfigTest(
+    AbstractPreconditionerConfigTest.PreconditionerConfigTest[
+        Type[EigenvalueCorrectedShampooPreconditionerConfig]
+    ]
+):
+    def _get_preconditioner_config_type(
+        self,
+    ) -> Type[EigenvalueCorrectedShampooPreconditionerConfig]:
+        return EigenvalueCorrectedShampooPreconditionerConfig

distributed_shampoo/utils/shampoo_preconditioner_list.py

@@ -627,6 +627,45 @@ def _check_factor_matrix_for_diagonality_nan_and_inf(
                 f"{factor_matrix.isinf().any()=}, {factor_matrix.isnan().any()=}."
             )
 
+    def _raise_exception_if_failure_tolerance_exceeded(
+        self,
+        success_tracker: list[bool],
+        preconditioner_index: int,
+        exception: Exception,
+    ) -> None:
+        """Raises an exception if the number of failed amortized computations exceeds the tolerance.
+
+        Resets the counter at the given index when all amortized computations are successful.
+
+        Args:
+            success_tracker (list[bool]): A list of booleans indicating whether the amortized computation was successful.
+            preconditioner_index (int): The index of the preconditioner.
+            exception (Exception): The exception to raise.
+
+        Raises:
+            exception (Exception): The exception to raise.
+
+        """
+        if all(success_tracker):
+            # Reset counter for failed amortized computations.
+            self._masked_failed_amortized_computation_counter_list[
+                preconditioner_index
+            ] = 0
+        else:
+            # Increment counter for failed amortized computations.
+            self._masked_failed_amortized_computation_counter_list[
+                preconditioner_index
+            ] += 1
+            # Raise the exception if the tolerance at the given index is exceeded.
+            failure_counter = self._masked_failed_amortized_computation_counter_list[
+                preconditioner_index
+            ]
+            tolerance = (
+                self._preconditioner_config.num_tolerated_failed_amortized_computations
+            )
+            if failure_counter > tolerance:
+                raise exception
+
     def update_preconditioners(
         self,
         masked_grad_list: tuple[Tensor, ...],
@@ -678,10 +717,16 @@ def _initialize_state_lists(
             self._inv_root_override,
             self._local_order_list,
         )
+        self._local_failed_amortized_computation_counter_list: list[int] = [0] * len(
+            self._local_kronecker_factors_list
+        )
 
         # Masked lists are the list of active preconditioners or values after filtering out gradients with None.
         self._masked_order_list: tuple[int, ...] = self._local_order_list
         self._masked_root_list: tuple[int, ...] = self._local_root_list
+        self._masked_failed_amortized_computation_counter_list: list[int] = (
+            self._local_failed_amortized_computation_counter_list
+        )
         self._masked_kronecker_factors_list: tuple[
             ShampooKroneckerFactorsListType,
             ...,
@@ -714,6 +759,14 @@ def compress_preconditioner_list(
             self._masked_root_list: tuple[int, ...] = compress_list(  # type: ignore[no-redef]
                 self._local_root_list, local_grad_selector
             )
+            self._masked_failed_amortized_computation_counter_list: list[int] = (  # type: ignore[no-redef]
+                list(
+                    compress_list(
+                        self._local_failed_amortized_computation_counter_list,
+                        local_grad_selector,
+                    )
+                )
+            )
             self._masked_kronecker_factors_list: tuple[  # type: ignore[no-redef]
                 ShampooKroneckerFactorsListType,
                 ...,
@@ -850,11 +903,14 @@ def _amortized_computation(self) -> None:
         with profiler.record_function(
             f"## {self.__class__.__name__}:{self._amortized_computation.__name__} ##"
         ):
-            for kronecker_factors, root in zip(
-                self._masked_kronecker_factors_list,
-                self._masked_root_list,
-                strict=True,
+            for idx, (kronecker_factors, root) in enumerate(
+                zip(
+                    self._masked_kronecker_factors_list,
+                    self._masked_root_list,
+                    strict=True,
+                )
             ):
+                success_tracker: list[bool] = []
                 for (
                     factor_matrix,
                     inv_factor_matrix,
@@ -898,8 +954,11 @@ def _amortized_computation(self) -> None:
                             epsilon=self._epsilon,
                             is_diagonal=bool(is_factor_matrix_diagonal),
                         ).to(dtype=inv_factor_matrix.dtype)
+                        # Add success to success tracker.
+                        success_tracker.append(True)
                     except Exception as exception:
-                        # Reuse previous matrix if matrix inverse root computation fails.
+                        # Add failure to success tracker.
+                        success_tracker.append(False)
                         logger.warning(
                             f"Matrix computation failed for factor matrix {factor_matrix_index} "
                             f"with {exception=}. Using previous inverted factor matrix and continuing..."
@@ -919,6 +978,15 @@ def _amortized_computation(self) -> None:
                         )
                     inv_factor_matrix.copy_(computed_inv_factor_matrix)
 
+                # Only reuse previous inverse roots if tolerance is not exceeded.
+                self._raise_exception_if_failure_tolerance_exceeded(
+                    success_tracker=success_tracker,
+                    preconditioner_index=idx,
+                    exception=ValueError(
+                        f"Exceeded tolerance for number of failed inverse root computations for {kronecker_factors.factor_matrix_indices}."
+                    ),
+                )
+
 
 class EigenvalueCorrectedShampooPreconditionerList(
     BaseShampooPreconditionerList[EigenvalueCorrectedShampooKroneckerFactorsList]
@@ -1098,7 +1166,10 @@ def _amortized_computation(self) -> None:
         with profiler.record_function(
             f"## {self.__class__.__name__}:{self._amortized_computation.__name__} ##"
         ):
-            for kronecker_factors in self._masked_kronecker_factors_list:
+            for idx, kronecker_factors in enumerate(
+                self._masked_kronecker_factors_list
+            ):
+                success_tracker: list[bool] = []
                 for (
                     factor_matrix,
                     factor_matrix_eigenvectors,
@@ -1129,8 +1200,11 @@ def _amortized_computation(self) -> None:
                             eigenvector_computation_config=eigenvector_computation_config,
                             is_diagonal=bool(is_factor_matrix_diagonal),
                         )
+                        # Add success to success tracker.
+                        success_tracker.append(True)
                     except Exception as exception:
-                        # Reuse previous matrix if matrix eigenvector computation fails.
+                        # Add failure to success tracker.
+                        success_tracker.append(False)
                         logger.warning(
                             f"Matrix computation failed for factor matrix {factor_matrix_index} "
                             f"with {exception=}. Using previous factor matrix eigenvectors and continuing..."
@@ -1149,3 +1223,12 @@ def _amortized_computation(self) -> None:
                             f"To mitigate, check factor matrix before the matrix computation: {factor_matrix=}"
                         )
                     factor_matrix_eigenvectors.copy_(computed_eigenvectors)
+
+                # Only reuse previous eigenvectors if tolerance is not exceeded.
+                self._raise_exception_if_failure_tolerance_exceeded(
+                    success_tracker=success_tracker,
+                    preconditioner_index=idx,
+                    exception=ValueError(
+                        f"Exceeded tolerance for number of failed eigenvector computations for {kronecker_factors.factor_matrix_indices}."
+                    ),
+                )