Using EmbeddingPlannerBase for lp planner

torchrec/distributed/planner/__init__.py

@@ -21,6 +21,9 @@
     - automatically building and selecting an optimized sharding plan.
 """
 
-from torchrec.distributed.planner.planners import EmbeddingShardingPlanner  # noqa
+from torchrec.distributed.planner.planners import (  # noqa  # noqa
+    EmbeddingPlannerBase,
+    EmbeddingShardingPlanner,
+)
 from torchrec.distributed.planner.types import ParameterConstraints, Topology  # noqa
 from torchrec.distributed.planner.utils import bytes_to_gb, sharder_name  # noqa

torchrec/distributed/planner/planners.py

@@ -251,6 +251,22 @@ def collective_plan(
             sharders,
         )
 
+    def hash_planner_context_inputs(self) -> int:
+        """
+        Generates a hash for all planner inputs except for partitioner, proposer, performance model, and stats.
+        These are all the inputs needed to verify whether a previously generated sharding plan is still valid in a new context.
+
+        Returns:
+            Generates a hash capturing topology, batch size, enumerator, storage reservation, stats and constraints.
+        """
+        return hash_planner_context_inputs(
+            self._topology,
+            self._batch_size,
+            self._enumerator,
+            self._storage_reservation,
+            self._constraints,
+        )
+
 
 class EmbeddingShardingPlanner(EmbeddingPlannerBase):
     """
@@ -368,22 +384,6 @@ def collective_plan(
             sharders,
         )
 
-    def hash_planner_context_inputs(self) -> int:
-        """
-        Generates a hash for all planner inputs except for partitioner, proposer, performance model, and stats.
-        These are all the inputs needed to verify whether a previously generated sharding plan is still valid in a new context.
-
-        Returns:
-            Generates a hash capturing topology, batch size, enumerator, storage reservation, stats and constraints.
-        """
-        return hash_planner_context_inputs(
-            self._topology,
-            self._batch_size,
-            self._enumerator,
-            self._storage_reservation,
-            self._constraints,
-        )
-
     def plan(
         self,
         module: nn.Module,