add warmup step

Author: samsja

distributed_shampoo/utils/shampoo_preconditioner_list.py

@@ -532,7 +532,7 @@ def _get_inverse_roots_from_override_with_high_order_default(
             )
 
     @abstractmethod
-    def _amortized_computation(self) -> None:
+    def _amortized_computation(self, step: int) -> None:
         """
         Computes the amortized computation needed for each Shampoo preconditioner implementation.
         This amortized computation is computation heavy work that cannot be done for each step.
@@ -631,7 +631,7 @@ def update_preconditioners(
             # In Shampoo, this is equivalent to computing the inverse factor matrix.
             # In Eigenvalue-Corrected Shampoo, this is equivalent to computing the eigenvector of the factor matrix.
             if perform_amortized_computation:
-                self._amortized_computation()
+                self._amortized_computation(step=step)
 
     def _initialize_state_lists(
         self,
@@ -797,7 +797,7 @@ def precondition(self, masked_grad_list: tuple[Tensor, ...]) -> tuple[Tensor, ..
             )
 
     @torch.compiler.disable
-    def _amortized_computation(self) -> None:
+    def _amortized_computation(self, step: int) -> None:
         # NOTE: This function currently only computes the matrix root inverse based on
         # the masked lists which combines both selection based on the distributor and where
         # grad is not None. Implicitly, this assumes that there are no changes between the
@@ -1032,7 +1032,7 @@ def precondition(self, masked_grad_list: tuple[Tensor, ...]) -> tuple[Tensor, ..
             return tuple(preconditioned_grad_list)
 
     @torch.compiler.disable
-    def _amortized_computation(self) -> None:
+    def _amortized_computation(self, step: int) -> None:
         # NOTE: This function currently only computes the preconditioner eigenvectors based on
         # the masked lists which combines both selection based on the distributor and where
         # grad is not None. Implicitly, this assumes that there are no changes between the
@@ -1071,6 +1071,7 @@ def _amortized_computation(self) -> None:
                             eigenvectors_estimate=factor_matrix_eigenvectors,
                             eigenvector_computation_config=eigenvector_computation_config,
                             is_diagonal=bool(is_factor_matrix_diagonal),
+                            step=step,
                         )
                         # Add success to success tracker.
                         success_tracker.append(True)

matrix_functions.py

@@ -624,6 +624,7 @@ def matrix_eigenvectors(
     eigenvectors_estimate: Tensor | None = None,
     eigenvector_computation_config: EigenvectorConfig = DefaultEighEigenvectorConfig,
     is_diagonal: bool = False,
+    step: int | None = None,
 ) -> Tensor:
     """Compute eigenvectors of matrix using eigendecomposition of symmetric positive (semi-)definite matrix.
             A = Q L Q^T => Q
@@ -668,9 +669,13 @@ def matrix_eigenvectors(
             retry_double_precision=eigenvector_computation_config.retry_double_precision,
         )
 
+        if step is None:
+            raise ValueError("step param is required when using EighEigenvectorConfig.")
+
         if (
             isinstance(eigenvector_computation_config, TopKCompressionEigenvectorConfig)
             and eigenvalues.shape[0] > eigenvector_computation_config.min_dim
+            and step > eigenvector_computation_config.warmup_steps
         ):
             effective_rank = compute_effective_rank(eigenvalues, eigenvector_computation_config.compression_t)
 

matrix_functions_types.py

@@ -151,6 +151,8 @@ class TopKCompressionEigenvectorConfig(EighEigenvectorConfig):
 
     compression_t: float = 0.95
 
+    warmup_steps: int = 0
+
     def __post_init__(self):
         if isinstance(self.topk_compression, float):
             if not 0 < self.topk_compression <= 1: