Skip to content

Commit 4a1b042

Browse files
aliafzalfacebook-github-bot
authored andcommitted
Add ConfigeratorStats to store sharding plan in config store (#3327)
Summary: internal 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. **How is a sharding plan stored in Configerator?** The Thrift definition of a sharding plan includes two fields: Topology and Dict[int, ShardingOption]. 1. Topology: The Topology field contains the information mentioned in this diff D79142495. 2. Dict[int, ShardingOption]: This field represents a dictionary where each key is a 64-bit hash of a sharding option, and the value is the corresponding Thrift-converted sharding option. The hash is calculated using the storage_hash function within the ShardingOption object, which takes into account factors such as the fqn, sharding type, and compute kernel. **How can a loaded plan be merged with an enumerated search space?** **Background:** When a plan is preserved during the logging stage, a hash is generated to ensure that the same plan can be loaded and validated later. The [hash is calculated](https://www.internalfb.com/code/fbsource/[fdf90ff2be9041f867bc6c9e4aec6ee94862fa11]/fbcode/torchrec/distributed/planner/types.py?lines=1010-1026) using input fields such as topology, batch size, constraints, storage reservation, and storage reservation policy, as well as fields from the sharding options like fqn, sharding type, kernel type, shards, and cache parameters. Once the plan is loaded and validated, we can safely assume that all loaded sharding options are a 1:1 map of enumerated sharded options. During the loading process, we traverse the enumerated search space, calculate the storage hash for each sharding option, look up the corresponding sharding option from the loaded plan, and replace the Shards of the enumerated sharding option with those of the loaded sharding option. This approach enables us to generate precise sharding options that can be seamlessly converted into a sharing plan as done by the planner and this also ensures consistent logging while also facilitating plan replay. Differential Revision: D81185992
1 parent 3371465 commit 4a1b042

File tree

1 file changed

+11
-0
lines changed

1 file changed

+11
-0
lines changed

torchrec/distributed/planner/types.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -635,6 +635,17 @@ 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=7).digest()
646+
hash_int = int.from_bytes(hash_bytes, byteorder="big")
647+
return hash_int
648+
638649
def __deepcopy__(
639650
self, memo: Optional[Dict[int, "ShardingOption"]]
640651
) -> "ShardingOption":

0 commit comments

Comments
 (0)