add egeinvector stats

Author: samsja

distributed_shampoo/distributed_shampoo.py

@@ -1250,3 +1250,27 @@ def load_distributed_state_dict(
                 param_group_to_load = param_groups_to_load[param_group_key]
                 for key, value in param_group_to_load.items():
                     group[key] = deepcopy(value)
+
+    
+    @torch.no_grad()
+    def eigenvector_stats(self, key_to_param: Iterator[tuple[str, torch.Tensor]], summary: bool = False):
+        # Create mapping from parameter to its name
+        param_to_key = {param: key for key, param in key_to_param}
+        
+        stats = {}
+        for idx, (state_lists, group) in enumerate(zip(self._per_group_state_lists, self.param_groups)):
+            shampoo_preconditioner_list = state_lists[SHAMPOO_PRECONDITIONER_LIST]
+            if isinstance(shampoo_preconditioner_list, EigenvalueCorrectedShampooPreconditionerList):
+                # Get eigenvalue stats for this group
+                group_eigen_stats = shampoo_preconditioner_list.eigenvector_stats()
+                
+                # Map each parameter to its stats
+                param_stats = {}
+                for param, eigen_stat in zip(group[PARAMS], group_eigen_stats):
+                    if param in param_to_key:
+                        param_key = param_to_key[param]
+                        param_stats[param_key] = eigen_stat
+                
+                stats[f"group_{idx}"] = param_stats
+           
+        return stats

distributed_shampoo/utils/shampoo_preconditioner_list.py

@@ -1093,8 +1093,7 @@ 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
+                    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(
@@ -1104,3 +1103,6 @@ def _amortized_computation(self, step: int) -> None:
                         f"The number of failed eigenvector computations for factors {kronecker_factors.factor_matrix_indices} exceeded the allowed tolerance."
                     ),
                 )
+                
+    def eigenvector_stats(self) -> tuple[EigenStats | None, ...]:
+        return tuple(kronecker_factors.eigen_stats for kronecker_factors in self._masked_kronecker_factors_list)

matrix_functions.py

@@ -633,6 +633,13 @@ def compression_ratio(self):
     def __repr__(self):
         return f"Effective rank: {self.effective_rank}, og_rank: {self.og_rank}, compression_ratio: {self.compression_ratio}"
 
+    def log_stats(self) -> dict[str, int|float]:
+        return {
+            "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,
@@ -704,7 +711,6 @@ def matrix_eigenvectors(
 
             if eigenvector_computation_config.auto:
                 topk = eigen_stats.effective_rank
-                print(eigen_stats)
             elif isinstance(eigenvector_computation_config.topk_compression, int):
                 topk = eigenvector_computation_config.topk_compression
             else: