add egein stats

Author: samsja

distributed_shampoo/utils/shampoo_preconditioner_list.py

@@ -24,7 +24,7 @@
 )
 from distributed_shampoo.utils.shampoo_block_info import BlockInfo
 from distributed_shampoo.utils.shampoo_utils import compress_list, get_dtype_size
-from matrix_functions import check_diagonal, matrix_eigenvectors, matrix_inverse_root
+from matrix_functions import EigenStats, check_diagonal, matrix_eigenvectors, matrix_inverse_root
 
 from matrix_functions_types import EigenvectorConfig, RootInvConfig
 from optimizer_modules import OptimizerModule
@@ -303,6 +303,7 @@ class EigenvalueCorrectedShampooKroneckerFactorsList(BaseShampooKroneckerFactors
 
     factor_matrices_eigenvectors: tuple[Tensor, ...]
     corrected_eigenvalues: Tensor
+    eigen_stats: EigenStats | None = None
 
     def __post_init__(self) -> None:
         super().__post_init__()
@@ -1065,8 +1066,7 @@ def _amortized_computation(self, step: int) -> None:
                         self._preconditioner_config.amortized_computation_config,
                     )
                     try:
-                        logger.info(f"TYPEEEE: {type(eigenvector_computation_config)}...")
-                        computed_eigenvectors = matrix_eigenvectors(
+                        computed_eigenvectors, eigen_stats = matrix_eigenvectors(
                             A=factor_matrix,
                             eigenvectors_estimate=factor_matrix_eigenvectors,
                             eigenvector_computation_config=eigenvector_computation_config,
@@ -1093,6 +1093,8 @@ def _amortized_computation(self, step: int) -> None:
                             f"To mitigate, check factor matrix before the matrix computation: {factor_matrix=}"
                         )
                     factor_matrix_eigenvectors.copy_(computed_eigenvectors)
+                    print(type(self._masked_kronecker_factors_list[idx]))
+                    # self._masked_kronecker_factors_list[idx].eigen_stats = eigen_stats
 
                 # Only reuse previous eigenvectors if tolerance is not exceeded.
                 self._raise_exception_if_failure_tolerance_exceeded(

matrix_functions.py

@@ -30,10 +30,13 @@
 )
 
 from torch import Tensor
+from typing import NamedTuple
+
 
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+
 class NewtonConvergenceFlag(enum.Enum):
     """
     Enum class for the state of the Newton / higher-order iteration method.
@@ -619,13 +622,24 @@ def compute_effective_rank(eigenvalues: torch.Tensor, threshold: float = 0.95) -
     return effective_rank
 
 
+class EigenStats(NamedTuple):
+    effective_rank: int
+    og_rank: int
+    
+    @property
+    def compression_ratio(self):
+        return 1 - self.effective_rank / self.og_rank
+    
+    def __repr__(self):
+        return f"Effective rank: {self.effective_rank}, og_rank: {self.og_rank}, compression_ratio: {self.compression_ratio}"
+
 def matrix_eigenvectors(
     A: Tensor,
     eigenvectors_estimate: Tensor | None = None,
     eigenvector_computation_config: EigenvectorConfig = DefaultEighEigenvectorConfig,
     is_diagonal: bool = False,
     step: int | None = None,
-) -> Tensor:
+) -> Tensor | EigenStats:
     """Compute eigenvectors of matrix using eigendecomposition of symmetric positive (semi-)definite matrix.
             A = Q L Q^T => Q
 
@@ -663,12 +677,20 @@ def matrix_eigenvectors(
             device=A.device,
         )
 
+    
     if isinstance(eigenvector_computation_config, EighEigenvectorConfig):
         eigenvalues, eigenvectors = matrix_eigenvalue_decomposition(
             A,
             retry_double_precision=eigenvector_computation_config.retry_double_precision,
         )
 
+        compression_t = eigenvector_computation_config.compression_t if isinstance(eigenvector_computation_config, TopKCompressionEigenvectorConfig) else TopKCompressionEigenvectorConfig().compression_t
+        
+        eigen_stats = EigenStats(
+            effective_rank=compute_effective_rank(eigenvalues, compression_t),
+            og_rank=eigenvalues.shape[0],
+        )
+        
         if step is None:
             raise ValueError("step param is required when using EighEigenvectorConfig.")
 
@@ -677,36 +699,19 @@ def matrix_eigenvectors(
             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)
 
-            # rank = int(os.environ.get("RANK", 0))
-            potential_compression_ratio = 1 - effective_rank / eigenvalues.shape[0]
 
-            # if rank == 0:
-            # import wandb
-
-            # wandb.log(
-            #     {
-            #         "effective_rank": effective_rank,
-            #         "og_rank": eigenvalues.shape[0],
-            #         "potential_compression_ratio": 1 - effective_rank / eigenvalues.shape[0],
-            #     }
-            # )
 
             if eigenvector_computation_config.auto:
-                topk = effective_rank
-                print(
-                    f"Effective rank: {effective_rank}, og_rank: {eigenvalues.shape[0]}, compression_ratio: {potential_compression_ratio}"
-                )
+                topk = eigen_stats.effective_rank
+                print(eigen_stats)
             elif isinstance(eigenvector_computation_config.topk_compression, int):
                 topk = eigenvector_computation_config.topk_compression
             else:
                 topk = int(eigenvector_computation_config.topk_compression * eigenvalues.shape[0])
 
-            if potential_compression_ratio < eigenvector_computation_config.min_compression_ratio:
-                print(
-                    f"Skipping eigenvector computation due to low compression ratio: {potential_compression_ratio}, effective_rank = {effective_rank}, og_rank = {eigenvalues.shape[0]}"
-                )
+            if eigen_stats.compression_ratio < eigenvector_computation_config.min_compression_ratio:
+                print(f"Skipping eigenvector computation due to low compression ratio: {eigen_stats}")
                 return eigenvectors
             # Sort eigenvalues and eigenvectors in descending order
             eigenvalues, indices = torch.sort(
@@ -720,9 +725,12 @@ def matrix_eigenvectors(
             mask[:, :topk] = 1.0
             eigenvectors = eigenvectors * mask
 
-        return eigenvectors
+        return eigenvectors, eigen_stats
 
     elif isinstance(eigenvector_computation_config, QRConfig):
+        
+        raise NotImplementedError("QRConfig is not implemented yet.")
+        
         assert eigenvectors_estimate is not None, "Estimate of eigenvectors is required when using QRConfig."
 
         eigenvectors = _compute_orthogonal_iterations(