Add cache params to hash

Summary: The reason for doing that is we want to distinguish sharding options with a different cache load factor. Differential Revision: D50523588
henrylhtsang · Oct 31, 2023 · 16b52ee · 16b52ee
1 parent 9434f25
commit 16b52ee
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/torchrec/distributed/planner/types.py b/torchrec/distributed/planner/types.py
@@ -338,6 +338,7 @@ def __hash__(self) -> int:
                 self.sharding_type,
                 self.compute_kernel,
                 tuple(self.shards),
+                self.cache_params,
             )
         )
 

diff --git a/torchrec/distributed/types.py b/torchrec/distributed/types.py
@@ -481,6 +481,17 @@ class CacheParams:
     precision: Optional[DataType] = None
     prefetch_pipeline: Optional[bool] = None
 
+    def __hash__(self) -> int:
+        return hash(
+            (
+                self.algorithm,
+                self.load_factor,
+                self.reserved_memory,
+                self.precision,
+                self.prefetch_pipeline,
+            )
+        )
+
 
 @dataclass
 class ParameterSharding: