Add custom dtype supports for subclasses of ShampooPreconditionerConfig (#249)

Summary:
Pull Request resolved: #249

1. Add `inv_factor_matrix_dtype` option in `RootInvShampooPreconditionerConfig`.
2. Add `factor_matrix_eigenvectors_dtype` and `factor_matrix_eigenvalues_dtype` options in `EigendecomposedShampooPreconditionerConfig`.
3. Add `factor_matrix_eigenvectors_dtype` and `corrected_eigenvalues_dtype` options in `EigenvalueCorrectedShampooPreconditionerConfig`.

Reviewed By: runame

Differential Revision: D82239916

fbshipit-source-id: 6848168df817eff044b6ca4042477a5392e16e63

Author: tsunghsienlee

distributed_shampoo/preconditioner/shampoo_preconditioner_list.py

@@ -36,7 +36,10 @@
 )
 from distributed_shampoo.shampoo_types import (
     AmortizedPreconditionerConfig,
+    EigendecomposedShampooPreconditionerConfig,
+    EigenvalueCorrectedShampooPreconditionerConfig,
     PreconditionerValueError,
+    RootInvShampooPreconditionerConfig,
 )
 from distributed_shampoo.utils.dict_zip_iterator import DictZipIterator
 from distributed_shampoo.utils.optimizer_modules import OptimizerModule
@@ -320,31 +323,31 @@ def from_block(cls, **kwargs: Any) -> "RootInvShampooKroneckerFactorsState":
 
         Args:
             block_info (BlockInfo): Information about the block, including methods to allocate tensors.
-            factor_matrix_dtype (torch.dtype): Data type for the factor matrices.
+            preconditioner_config (RootInvShampooPreconditionerConfig): Configuration for the preconditioner.
             preconditioned_dims (tuple[int, ...]): Dimensions for which the factor matrices are preconditioned.
-            block_dtype (torch.dtype): Data type for the block.
 
         Returns:
             kronecker_factors_state (RootInvShampooKroneckerFactorsState): An instance of RootInvShampooKroneckerFactorsState with initialized inverse factor matrices.
         """
         block_info: BlockInfo = kwargs["block_info"]
-        factor_matrix_dtype: torch.dtype = kwargs["factor_matrix_dtype"]
+        preconditioner_config: RootInvShampooPreconditionerConfig = kwargs[
+            "preconditioner_config"
+        ]
         preconditioned_dims: tuple[int, ...] = kwargs["preconditioned_dims"]
-        block_dtype: torch.dtype = kwargs["block_dtype"]
 
         return cls(
             **asdict(
                 BaseShampooKroneckerFactorsState.from_block(
                     block_info=block_info,
-                    factor_matrix_dtype=factor_matrix_dtype,
+                    factor_matrix_dtype=preconditioner_config.factor_matrix_dtype,
                     preconditioned_dims=preconditioned_dims,
                 )
             ),
             # Initialize inv_factor_matrices as identity matrices.
             inv_factor_matrices=tuple(
                 block_info.allocate_eye_tensor(
                     n=dim,
-                    dtype=block_dtype,
+                    dtype=preconditioner_config.inv_factor_matrix_dtype,
                     device=block_info.param.device,
                 )
                 for dim in preconditioned_dims
@@ -517,31 +520,31 @@ def from_block(cls, **kwargs: Any) -> "EigendecomposedShampooKroneckerFactorsSta
 
         Args:
             block_info (BlockInfo): Information about the block, including methods to allocate tensors.
-            factor_matrix_dtype (torch.dtype): Data type for the factor matrices.
+            preconditioner_config (EigendecomposedShampooPreconditionerConfig): Configuration for the preconditioner.
             preconditioned_dims (tuple[int, ...]): Dimensions for which the factor matrices are preconditioned.
-            block_dtype (torch.dtype): Data type for the block.
 
         Returns:
             kronecker_factors_state (EigendecomposedShampooKroneckerFactorsState): An instance of EigendecomposedShampooKroneckerFactorsState.
         """
         block_info: BlockInfo = kwargs["block_info"]
-        factor_matrix_dtype: torch.dtype = kwargs["factor_matrix_dtype"]
+        preconditioner_config: EigendecomposedShampooPreconditionerConfig = kwargs[
+            "preconditioner_config"
+        ]
         preconditioned_dims: tuple[int, ...] = kwargs["preconditioned_dims"]
-        block_dtype: torch.dtype = kwargs["block_dtype"]
 
         return cls(
             **asdict(
                 BaseShampooKroneckerFactorsState.from_block(
                     block_info=block_info,
-                    factor_matrix_dtype=factor_matrix_dtype,
+                    factor_matrix_dtype=preconditioner_config.factor_matrix_dtype,
                     preconditioned_dims=preconditioned_dims,
                 )
             ),
             # Initialize factor_matrices_eigenvectors as identity matrices.
             factor_matrices_eigenvectors=tuple(
                 block_info.allocate_eye_tensor(
                     n=dim,
-                    dtype=block_dtype,
+                    dtype=preconditioner_config.factor_matrix_eigenvectors_dtype,
                     device=block_info.param.device,
                 )
                 for dim in preconditioned_dims
@@ -550,7 +553,7 @@ def from_block(cls, **kwargs: Any) -> "EigendecomposedShampooKroneckerFactorsSta
             factor_matrices_eigenvalues=tuple(
                 block_info.allocate_ones_tensor(
                     size=(dim,),
-                    dtype=block_dtype,
+                    dtype=preconditioner_config.factor_matrix_eigenvalues_dtype,
                     device=block_info.param.device,
                 )
                 for dim in preconditioned_dims
@@ -760,41 +763,41 @@ def from_block(
 
         Args:
             block_info (BlockInfo): Information about the block, including methods to allocate tensors.
-            factor_matrix_dtype (torch.dtype): Data type for the factor matrices.
+            preconditioner_config (EigenvalueCorrectedShampooPreconditionerConfig): Configuration for the preconditioner.
             preconditioned_dims (tuple[int, ...]): Dimensions for which the factor matrices are preconditioned.
-            block_dtype (torch.dtype): Data type for the block.
             dims (tuple[int, ...]): Dimensions of the block.
 
         Returns:
             kronecker_factors_state (EigenvalueCorrectedShampooKroneckerFactorsState): An instance of EigenvalueCorrectedShampooKroneckerFactorsState.
         """
         block_info: BlockInfo = kwargs["block_info"]
-        factor_matrix_dtype: torch.dtype = kwargs["factor_matrix_dtype"]
+        preconditioner_config: EigenvalueCorrectedShampooPreconditionerConfig = kwargs[
+            "preconditioner_config"
+        ]
         preconditioned_dims: tuple[int, ...] = kwargs["preconditioned_dims"]
-        block_dtype: torch.dtype = kwargs["block_dtype"]
         dims: tuple[int, ...] = kwargs["dims"]
 
         return EigenvalueCorrectedShampooKroneckerFactorsState(
             **asdict(
                 BaseShampooKroneckerFactorsState.from_block(
                     block_info=block_info,
-                    factor_matrix_dtype=factor_matrix_dtype,
+                    factor_matrix_dtype=preconditioner_config.factor_matrix_dtype,
                     preconditioned_dims=preconditioned_dims,
                 )
             ),
             # Initialize factor_matrices_eigenvectors as identity matrices.
             factor_matrices_eigenvectors=tuple(
                 block_info.allocate_eye_tensor(
                     n=dim,
-                    dtype=block_dtype,
+                    dtype=preconditioner_config.factor_matrix_eigenvectors_dtype,
                     device=block_info.param.device,
                 )
                 for dim in preconditioned_dims
             ),
             corrected_eigenvalues=block_info.allocate_zeros_tensor(
                 # Note that the corrected eigenvalues are not affected by the preconditioned_dims.
                 size=tuple(dims),
-                dtype=block_dtype,
+                dtype=preconditioner_config.corrected_eigenvalues_dtype,
                 device=block_info.param.device,
             ),
         )
@@ -1133,9 +1136,8 @@ def _create_kronecker_factors_state(
             )
             block_state[SHAMPOO] = kronecker_factors_state_type.from_block(
                 block_info=block_info,
-                factor_matrix_dtype=self._preconditioner_config.factor_matrix_dtype,
+                preconditioner_config=self._preconditioner_config,
                 preconditioned_dims=preconditioned_dims,
-                block_dtype=block.dtype,
                 dims=dims,
             )
             kronecker_factors_unwrapped.append(
@@ -1316,17 +1318,30 @@ def _precondition_grad(
         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)})."
+
+        # Extract all dtypes and assert they are unique
+        assert (
+            len(unique_dtypes := {p.dtype for p in preconditioner_list}) <= 1
+        ), f"All preconditioners must have the same dtype, but found: {unique_dtypes}"
+
+        # 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(
-                grad, next(preconditioner_list_iter), dims=dims
+                # Use the single target dtype for all operations
+                grad.to(dtype=target_dtype),
+                # Use the actual iterator for the operation
+                next(preconditioner_list_iter),
+                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)
 
     @overload
     @staticmethod

distributed_shampoo/preconditioner/tests/shampoo_preconditioner_list_test.py

@@ -1033,7 +1033,11 @@ def _amortized_computation_properties(self) -> AmortizedComputationProperties:
 
     @property
     def _default_preconditioner_config(self) -> RootInvShampooPreconditionerConfig:
-        return replace(DefaultShampooConfig, factor_matrix_dtype=torch.float64)
+        return replace(
+            DefaultShampooConfig,
+            factor_matrix_dtype=torch.float64,
+            inv_factor_matrix_dtype=torch.float64,
+        )
 
     @property
     def _preconditioner_list_factory(self) -> Callable[..., PreconditionerList]:
@@ -1059,6 +1063,8 @@ def _default_preconditioner_config(  # type: ignore[override]
         return EigendecomposedShampooPreconditionerConfig(
             amortized_computation_config=QREigendecompositionConfig(),
             factor_matrix_dtype=torch.float64,
+            factor_matrix_eigenvectors_dtype=torch.float64,
+            factor_matrix_eigenvalues_dtype=torch.float64,
         )
 
     @property
@@ -1077,7 +1083,12 @@ def _amortized_computation_properties(self) -> AmortizedComputationProperties:
     def _default_preconditioner_config(
         self,
     ) -> EigenvalueCorrectedShampooPreconditionerConfig:
-        return replace(DefaultSOAPConfig, factor_matrix_dtype=torch.float64)
+        return replace(
+            DefaultSOAPConfig,
+            factor_matrix_dtype=torch.float64,
+            factor_matrix_eigenvectors_dtype=torch.float64,
+            corrected_eigenvalues_dtype=torch.float64,
+        )
 
     @property
     def _preconditioner_list_factory(self) -> Callable[..., PreconditionerList]:

distributed_shampoo/shampoo_types.py

@@ -298,13 +298,15 @@ class RootInvShampooPreconditionerConfig(ShampooPreconditionerConfig):
                                     |          (^0.1667), the default inverse exponent 1/(2*3) since inverse_exponent_override[3][2] is not specified
                                     |
                             no preconditioning since inverse_exponent_override[3][0]=0.0
+        inv_factor_matrix_dtype (torch.dtype): Data type for inverse factor matrix. (Default: torch.float32)
 
 
     """
 
     amortized_computation_config: RootInvConfig = field(
         default_factory=lambda: DefaultEigenConfig
     )
+    inv_factor_matrix_dtype: torch.dtype = torch.float32
 
 
 DefaultShampooConfig = RootInvShampooPreconditionerConfig()
@@ -353,13 +355,17 @@ class EigendecomposedShampooPreconditionerConfig(ShampooPreconditionerConfig):
                                     |          (^0.1667), the default inverse exponent 1/(2*3) since inverse_exponent_override[3][2] is not specified
                                     |
                             no preconditioning since inverse_exponent_override[3][0]=0.0
+        factor_matrix_eigenvectors_dtype (torch.dtype): Data type for factor matrix eigenvectors. (Default: torch.float32)
+        factor_matrix_eigenvalues_dtype (torch.dtype): Data type for factor matrix eigenvalues. (Default: torch.float32)
 
 
     """
 
     amortized_computation_config: EigendecompositionConfig = field(
         default_factory=lambda: DefaultEigendecompositionConfig
     )
+    factor_matrix_eigenvectors_dtype: torch.dtype = torch.float32
+    factor_matrix_eigenvalues_dtype: torch.dtype = torch.float32
 
 
 @dataclass(kw_only=True)
@@ -410,6 +416,8 @@ class EigenvalueCorrectedShampooPreconditionerConfig(AmortizedPreconditionerConf
         inverse_exponent_override (dict[int, float]): The inverse_exponent_override attribute is a dictionary that allows for customizing the inverse exponent used in eigenvalue correction.
             The keys of the dictionary represent the order of the tensor, and the values are the exponent override values. For example, if we want to use a custom inverse exponent for 3-D tensors, we can set inverse_exponent_override as inverse_exponent_override={3: 0.25}.
             Note that the inverse_exponent_override dictionary can contain multiple entries for different tensor orders. If the order of the tensor is not specified in the dictionary, the default exponent, 1/2, will be used. (Default: {})
+        factor_matrix_eigenvectors_dtype (torch.dtype): Data type for factor matrix eigenvectors. (Default: torch.float32)
+        corrected_eigenvalues_dtype (torch.dtype): Data type for corrected eigenvalues. (Default: torch.float32)
 
     """
 
@@ -418,6 +426,8 @@ class EigenvalueCorrectedShampooPreconditionerConfig(AmortizedPreconditionerConf
     )
     ignored_basis_change_dims: dict[int, list[int]] = field(default_factory=dict)
     inverse_exponent_override: dict[int, float] = field(default_factory=dict)
+    factor_matrix_eigenvectors_dtype: torch.dtype = torch.float32
+    corrected_eigenvalues_dtype: torch.dtype = torch.float32
 
     def __post_init__(self) -> None:
         super().__post_init__()