Add hashing for Topology

Differential Revision: D76004583

Author: Caner Gocmen

torchrec/distributed/planner/tests/test_types.py

@@ -18,6 +18,7 @@
     ParameterConstraints,
     Shard,
     ShardingOption,
+    Topology,
 )
 from torchrec.distributed.types import (
     BoundsCheckMode,
@@ -214,6 +215,54 @@ def test_module_pooled_mch_ec(self) -> None:
         self.assertEqual(sharding_option.is_pooled, False)
 
 
+class TestTopologyHash(unittest.TestCase):
+    def test_hash_equality(self) -> None:
+        # Create two identical Topology instances
+        topology1 = Topology(
+            world_size=2,
+            compute_device="cuda",
+            hbm_cap=1024 * 1024 * 2,
+            local_world_size=2,
+        )
+
+        topology2 = Topology(
+            world_size=2,
+            compute_device="cuda",
+            hbm_cap=1024 * 1024 * 2,
+            local_world_size=2,
+        )
+
+        # Verify that the hash values are equal
+        self.assertEqual(
+            hash(topology1),
+            hash(topology2),
+            "Hashes should be equal for identical Topology instances",
+        )
+
+    def test_hash_inequality(self) -> None:
+        # Create two different Topology instances
+        topology1 = Topology(
+            world_size=2,
+            compute_device="cuda",
+            hbm_cap=1024 * 1024 * 2,
+            local_world_size=2,
+        )
+
+        topology2 = Topology(
+            world_size=4,  # Different world_size
+            compute_device="cuda",
+            hbm_cap=1024 * 1024 * 2,
+            local_world_size=2,
+        )
+
+        # Verify that the hash values are different
+        self.assertNotEqual(
+            hash(topology1),
+            hash(topology2),
+            "Hashes should be different for different Topology instances",
+        )
+
+
 class TestParameterConstraintsHash(unittest.TestCase):
 
     def test_hash_equality(self) -> None:

torchrec/distributed/planner/types.py

@@ -248,6 +248,10 @@ def get_bw(
 
 
 class Topology:
+    """
+    Representation of a network of devices in a cluster.
+    """
+
     def __init__(
         self,
         world_size: int,
@@ -396,6 +400,40 @@ def __repr__(self) -> str:
         topology_repr += str(self._comms_bandwidths) + "\n"
         return topology_repr
 
+    def __hash__(self) -> int:
+        """
+        Compute a hash value for this Topology instance.
+
+        This allows Topology objects to be used as dictionary keys or in sets.
+        The hash is based on the key attributes that define the topology.
+
+        Returns:
+            int: A hash value for this Topology instance.
+        """
+
+        # Compute hbms and ddrs from the decives
+        hbms = [device.storage.hbm for device in self._devices]
+        ddrs = [device.storage.ddr for device in self._devices]
+
+        # Combine all attributes into a hashable tuple
+        hashable_tuple = (
+            self._world_size,
+            self._compute_device,
+            frozenset(hbms),
+            frozenset(ddrs),
+            self._local_world_size,
+            self._hbm_mem_bw,
+            self._ddr_mem_bw,
+            self._hbm_to_ddr_mem_bw,
+            self._comms_bandwidths.intra_host_bw,
+            self._comms_bandwidths.inter_host_bw,
+            self._bwd_compute_multiplier,
+            self._weighted_feature_bwd_compute_multiplier,
+            self._uneven_sharding_perf_multiplier,
+        )
+
+        return hash(hashable_tuple)
+
 
 # ---- INPUT / OUTPUT ----- #