Replace lambda function in matrix_function_types.py and shampoo_types.py to a private static method for pickable concern

Summary:
This is the same spirit of D82836543, which changes the default value of `scale_fn` of `SignDescentPreconditionerConfig` to a private static method for pickling concern.

We found a few other occurrences of `lambda` function in `shampoo_types.py` and `matrix_function_types.py`. In order for shampoo state_dict (in particular, `param_groups`) to be compatibale with pytorch, we need to replace all lambda functions with a normal function, because `torch.save()` uses pickle for serialization, and pickle cannot serialize lambda function but only regular function.

Reviewed By: anana10c, hjmshi

Differential Revision: D85902488

fbshipit-source-id: ead3637fe2202ed10681312bf0a0652036be32a9

Author: wz337

distributed_shampoo/distributed_shampoo.py

@@ -1466,6 +1466,17 @@ def _post_state_dict_hook(optimizer: Optimizer, state_dict: StateDict) -> None:
         Returns:
             None: The state_dict is modified in-place.
         """
+
+        def _has_lambda_recursively(obj: Any) -> bool:
+            """Recursively check if an object contains lambda functions."""
+            if isinstance(obj, LambdaType):
+                return True
+            if is_dataclass(obj):
+                return any(
+                    _has_lambda_recursively(getattr(obj, f.name)) for f in fields(obj)
+                )
+            return False
+
         # for state exist on the ranks
         state_dict["state"] = {
             k: extract_state_dict_content(v) for k, v in state_dict["state"].items()
@@ -1477,9 +1488,7 @@ def _post_state_dict_hook(optimizer: Optimizer, state_dict: StateDict) -> None:
         for group in state_dict["param_groups"]:
             param_ids.extend(group["params"])
             for v in group.values():
-                if is_dataclass(v) and any(
-                    isinstance(getattr(v, f.name), LambdaType) for f in fields(v)
-                ):
+                if _has_lambda_recursively(v):
                     logger.warning(
                         f"Found {v=}. Note that lambda function cannot be pickled. "
                         "torch.save() cannot serialize lambda functions, because it "

distributed_shampoo/preconditioner/matrix_functions_types.py

@@ -84,6 +84,10 @@ class EigendecompositionConfig(MatrixFunctionConfig):
     Moreover, we have ||B||_F = ||Q^T A Q||_F = ||A||_F.
     Hence, the two relative errors are also equivalent: ||A - A'||_F / ||A||_F = ||B - diag(B)||_F / ||B||_F.
 
+    Note: When using custom rank_deficient_stability_config, avoid lambda functions as they may cause
+    pickling issues during serialization/deserialization. Use regular named functions
+    instead for better compatibility with distributed training and checkpointing.
+
     Attributes:
         rank_deficient_stability_config (RankDeficientStabilityConfig): Configuration for handling/stabilizing rank-deficient matrices. (Default: DefaultPerturbationConfig)
             TODO: generalize this to MatrixFunctionConfig
@@ -92,8 +96,12 @@ class EigendecompositionConfig(MatrixFunctionConfig):
 
     """
 
+    @staticmethod
+    def _get_default_rank_deficient_stability_config() -> RankDeficientStabilityConfig:
+        return DefaultPerturbationConfig
+
     rank_deficient_stability_config: RankDeficientStabilityConfig = field(
-        default_factory=lambda: DefaultPerturbationConfig
+        default_factory=_get_default_rank_deficient_stability_config
     )
     tolerance: float = 0.0
 
@@ -238,13 +246,28 @@ class OrthogonalizationConfig(MatrixFunctionConfig):
 
     If the reduced SVD of the matrix A is given by A = U S V^T, then the orthogonalized/closest orthogonal matrix is U V^T.
 
+    Note: When using custom scale_by_dims_fn, avoid lambda functions as they may cause
+    pickling issues during serialization/deserialization. Use regular named functions
+    instead for better compatibility with distributed training and checkpointing.
+
     Attributes:
         scale_by_dims_fn (Callable[[int, int], float]): Function to scale the orthogonalized matrix by some function of the dimensions of the matrix.
-            (Default: lambda d_in, d_out: 1.0)
+            (Default: _default_scale_by_dims_fn)
 
     """
 
-    scale_by_dims_fn: Callable[[int, int], float] = lambda d_in, d_out: 1.0
+    @staticmethod
+    def _default_scale_by_dims_fn(d_in: int, d_out: int) -> float:
+        """Default scaling function that returns 1.0 (no scaling)."""
+        return 1.0
+
+    @staticmethod
+    def _get_default_scale_by_dims_fn() -> Callable[[int, int], float]:
+        return OrthogonalizationConfig._default_scale_by_dims_fn
+
+    scale_by_dims_fn: Callable[[int, int], float] = field(
+        default_factory=_get_default_scale_by_dims_fn
+    )
 
 
 @dataclass(kw_only=True)

distributed_shampoo/shampoo_types.py

@@ -259,6 +259,10 @@ def __post_init__(self) -> None:
 class RootInvShampooPreconditionerConfig(ShampooPreconditionerConfig):
     """Configuration for Shampoo preconditioner computation with caching of the root inverse factor matrices.
 
+    Note: When using custom amortized_computation_config, avoid lambda functions as they may cause
+    pickling issues during serialization/deserialization. Use regular named functions
+    instead for better compatibility with distributed training and checkpointing.
+
     Attributes:
         amortized_computation_config (RootInvConfig): Configuration for the inverse-root computation. (Default: DefaultEigenConfig)
         num_tolerated_failed_amortized_computations (int): Number of failed amortized computations to tolerate before raising an error. (Default: 3)
@@ -303,8 +307,12 @@ class RootInvShampooPreconditionerConfig(ShampooPreconditionerConfig):
 
     """
 
+    @staticmethod
+    def _get_default_amortized_computation_config() -> RootInvConfig:
+        return DefaultEigenConfig
+
     amortized_computation_config: RootInvConfig = field(
-        default_factory=lambda: DefaultEigenConfig
+        default_factory=_get_default_amortized_computation_config
     )
     inv_factor_matrix_dtype: torch.dtype = torch.float32
 
@@ -316,6 +324,10 @@ class RootInvShampooPreconditionerConfig(ShampooPreconditionerConfig):
 class EigendecomposedShampooPreconditionerConfig(ShampooPreconditionerConfig):
     """Configuration for Shampoo preconditioner computation with caching of the eigendecomposed factor matrices.
 
+    Note: When using custom amortized_computation_config, avoid lambda functions as they may cause
+    pickling issues during serialization/deserialization. Use regular named functions
+    instead for better compatibility with distributed training and checkpointing.
+
     Attributes:
         amortized_computation_config (EigendecompositionConfig): Configuration for the eigendecomposition computation. (Default: DefaultEigendecompositionConfig)
         num_tolerated_failed_amortized_computations (int): Number of failed amortized computations to tolerate before raising an error. (Default: 3)
@@ -361,8 +373,12 @@ class EigendecomposedShampooPreconditionerConfig(ShampooPreconditionerConfig):
 
     """
 
+    @staticmethod
+    def _get_default_amortized_computation_config() -> EigendecompositionConfig:
+        return DefaultEigendecompositionConfig
+
     amortized_computation_config: EigendecompositionConfig = field(
-        default_factory=lambda: DefaultEigendecompositionConfig
+        default_factory=_get_default_amortized_computation_config
     )
     factor_matrix_eigenvectors_dtype: torch.dtype = torch.float32
     factor_matrix_eigenvalues_dtype: torch.dtype = torch.float32
@@ -375,6 +391,10 @@ class EigenvalueCorrectedShampooPreconditionerConfig(AmortizedPreconditionerConf
     Recall that in eigenvalue-corrected Shampoo, the eigenvectors and eigenvalues of the factor matrices are computed separately and stored in place of the full inverted preconditioner, as opposed to the single inverse-root computation of the factor matrices in Shampoo.
     In eigenvalue-corrected Shampoo, the eigenvectors are updated periodically like the inverted preconditioners in Shampoo, but the eigenvalues are updated every iteration.
 
+    Note: When using custom amortized_computation_config, avoid lambda functions as they may cause
+    pickling issues during serialization/deserialization. Use regular named functions
+    instead for better compatibility with distributed training and checkpointing.
+
     Attributes:
         amortized_computation_config (EigendecompositionConfig): Configuration for the eigenvector computation.
             (Default: DefaultEigendecompositionConfig)
@@ -421,8 +441,12 @@ class EigenvalueCorrectedShampooPreconditionerConfig(AmortizedPreconditionerConf
 
     """
 
+    @staticmethod
+    def _get_default_amortized_computation_config() -> EigendecompositionConfig:
+        return DefaultEigendecompositionConfig
+
     amortized_computation_config: EigendecompositionConfig = field(
-        default_factory=lambda: DefaultEigendecompositionConfig
+        default_factory=_get_default_amortized_computation_config
     )
     ignored_basis_change_dims: dict[int, list[int]] = field(default_factory=dict)
     inverse_exponent_override: dict[int, float] = field(default_factory=dict)
@@ -497,14 +521,23 @@ class SpectralDescentPreconditionerConfig(PreconditionerConfig):
     Which parameters are reshaped to 2D is determined by the max_preconditioner_dim argument in DistributedShampoo.
     If all >2D parameters should be guaranteed to be reshaped to 2D, then max_preconditioner_dim=math.inf and distributed_config.target_parameter_dimensionality=2 has to be used.
 
+
+    Note: When using custom orthogonalization config, avoid lambda functions as they may cause
+    pickling issues during serialization/deserialization. Use regular named functions
+    instead for better compatibility with distributed training and checkpointing.
+
     Attributes:
         orthogonalization_config (OrthogonalizationConfig): Configuration for orthogonalization of the search direction.
             (Default: DefaultNewtonSchulzOrthogonalizationConfig)
 
     """
 
+    @staticmethod
+    def _default_orthogonalization_config() -> OrthogonalizationConfig:
+        return DefaultNewtonSchulzOrthogonalizationConfig
+
     orthogonalization_config: OrthogonalizationConfig = field(
-        default_factory=lambda: DefaultNewtonSchulzOrthogonalizationConfig
+        default_factory=_default_orthogonalization_config
     )
 
 
@@ -595,12 +628,20 @@ class LoadBalancingConfig:
     The `cost_model` defines how the cost of a tensor is computed, and the distributor uses this cost to partition workloads.
     By default, it uses `AlignedMemoryCostModel`, other options include `PolynomialComputationalCostModel`.
 
+    Note: When using custom cost_model, avoid lambda functions as they may cause
+    pickling issues during serialization/deserialization. Use regular named functions
+    instead for better compatibility with distributed training and checkpointing.
+
     Args:
         cost_model (CostModel): The cost model used for load balancing. (Default: DefaultCostModel)
 
     """
 
-    cost_model: CostModel = field(default_factory=lambda: DefaultCostModel)
+    @staticmethod
+    def _get_default_cost_model() -> CostModel:
+        return DefaultCostModel
+
+    cost_model: CostModel = field(default_factory=_get_default_cost_model)
 
 
 @dataclass(init=False)
@@ -659,6 +700,10 @@ class DDPDistributedConfig(DistributedConfig):
 
     Enables distributed computation and optimizer states (like ZeRO-1) via DTensor for Shampoo.
 
+    Note: When using custom load_balancing_config, avoid lambda functions as they may cause
+    pickling issues during serialization/deserialization. Use regular named functions
+    instead for better compatibility with distributed training and checkpointing.
+
     Attributes:
         target_parameter_dimensionality (int | float): The idealized parameter dimensionality for a given algorithm.
             The dimensions of parameters and gradients will be merged (after squeezing dimensions of size 1) while respecting max_preconditioner_dim until the tensor has target_parameter_dimensionality dimensions left.
@@ -679,8 +724,13 @@ class DDPDistributedConfig(DistributedConfig):
     communication_dtype: torch.dtype = torch.float32
     num_trainers_per_group: int = -1
     communicate_params: bool = False
+
+    @staticmethod
+    def _get_default_load_balancing_config() -> LoadBalancingConfig:
+        return LoadBalancingConfig()
+
     load_balancing_config: LoadBalancingConfig = field(
-        default_factory=lambda: LoadBalancingConfig()
+        default_factory=_get_default_load_balancing_config
     )
 
 

distributed_shampoo/tests/distributed_shampoo_test.py

@@ -15,13 +15,15 @@
 import re
 import unittest
 from collections.abc import Callable
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, field, replace
 from typing import Any, cast
 
 import torch
 from distributed_shampoo.distributed_shampoo import DistributedShampoo
 from distributed_shampoo.preconditioner.matrix_functions_types import (
+    DefaultNewtonSchulzOrthogonalizationConfig,
     EigenConfig,
+    OrthogonalizationConfig,
     PseudoInverseConfig,
 )
 from distributed_shampoo.shampoo_types import (
@@ -1159,17 +1161,17 @@ def test_state_dict_warning(self) -> None:
         self.assertCountEqual(osd.keys(), ["state", "param_groups"])
 
         @dataclass(kw_only=True)
-        class SpectralDescentPreconditionerConfigWithLambda(
-            SpectralDescentPreconditionerConfig
-        ):
+        class SpectralDescentPreconditionerConfigWithLambda(PreconditionerConfig):
             """
-            Creating a preconditioner config with a dummy lambda function to make sure the
+            Creating a orthogonalization config with a dummy lambda function to make sure the
             warning from `_post_state_dict_hook` emit.
             """
 
-            scale_fn: Callable[[Tensor], float | Tensor] = lambda grad: 1.0
+            orthogonalization_config: OrthogonalizationConfig = field(
+                default_factory=lambda: DefaultNewtonSchulzOrthogonalizationConfig
+            )
 
-        self._optimizer.param_groups[0]["preconditioner_config"] = (
+        self._optimizer.param_groups[0]["orthogonalization_config"] = (
             SpectralDescentPreconditionerConfigWithLambda()
         )
         logger = logging.getLogger("distributed_shampoo.distributed_shampoo")