Skip to content

Commit 24c8dfd

Browse files
aliafzalfacebook-github-bot
authored andcommitted
Add ConfigeratorStats to store sharding plan in config store (#3327)
Summary: Pull Request resolved: #3327 Context: This change is part of the effort in improving planners overall UX and reliability. This Diff: 1. Add ConfigeratorStats to upload sharding plan to config store. Differential Revision: D81185992
1 parent 527e8b8 commit 24c8dfd

File tree

1 file changed

+12
-0
lines changed

1 file changed

+12
-0
lines changed

torchrec/distributed/planner/types.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -635,6 +635,18 @@ def __hash__(self) -> int:
635635
)
636636
)
637637

638+
def storage_hash(self) -> int:
639+
"""
640+
Hash needed to preserve sharding option uniquely based on input before
641+
planning. This is needed to restore sharding option from the loaded plan.
642+
"""
643+
# Use BLAKE2b for deterministic hashing, constrained to 32-bit signed int range
644+
hash_str = f"{self.fqn}|{self.sharding_type}|{self.compute_kernel}"
645+
hash_bytes = hashlib.blake2b(hash_str.encode("utf-8"), digest_size=8).digest()
646+
hash_int = int.from_bytes(hash_bytes, byteorder="big")
647+
# Constrain to 32-bit signed integer range for Thrift compatibility
648+
return hash_int % 2147483647 # 2^31 - 1
649+
638650
def __deepcopy__(
639651
self, memo: Optional[Dict[int, "ShardingOption"]]
640652
) -> "ShardingOption":

0 commit comments

Comments
 (0)