Skip to content

Commit 7dc79b9

Browse files
aliafzalfacebook-github-bot
authored andcommitted
Moving hash function to base planner class (#3238)
Summary: Moving planner hash function into EmbeddingPlannerBase to give all planner implementations access to this feature Reviewed By: aporialiao Differential Revision: D78996456
1 parent 44f4bb5 commit 7dc79b9

File tree

1 file changed

+16
-16
lines changed

1 file changed

+16
-16
lines changed

torchrec/distributed/planner/planners.py

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,22 @@ def collective_plan(
251251
sharders,
252252
)
253253

254+
def hash_planner_context_inputs(self) -> int:
255+
"""
256+
Generates a hash for all planner inputs except for partitioner, proposer, performance model, and stats.
257+
These are all the inputs needed to verify whether a previously generated sharding plan is still valid in a new context.
258+
259+
Returns:
260+
Generates a hash capturing topology, batch size, enumerator, storage reservation, stats and constraints.
261+
"""
262+
return hash_planner_context_inputs(
263+
self._topology,
264+
self._batch_size,
265+
self._enumerator,
266+
self._storage_reservation,
267+
self._constraints,
268+
)
269+
254270

255271
class EmbeddingShardingPlanner(EmbeddingPlannerBase):
256272
"""
@@ -368,22 +384,6 @@ def collective_plan(
368384
sharders,
369385
)
370386

371-
def hash_planner_context_inputs(self) -> int:
372-
"""
373-
Generates a hash for all planner inputs except for partitioner, proposer, performance model, and stats.
374-
These are all the inputs needed to verify whether a previously generated sharding plan is still valid in a new context.
375-
376-
Returns:
377-
Generates a hash capturing topology, batch size, enumerator, storage reservation, stats and constraints.
378-
"""
379-
return hash_planner_context_inputs(
380-
self._topology,
381-
self._batch_size,
382-
self._enumerator,
383-
self._storage_reservation,
384-
self._constraints,
385-
)
386-
387387
def plan(
388388
self,
389389
module: nn.Module,

0 commit comments

Comments
 (0)