Update DeltaStore API (#3469)

Summary:
Pull Request resolved: #3469

internal

General Context: We are in the process of transition to a unified DeltaTracker and this is 2/n diffs representing  changes towards the transition.

Specific Context: Update DeltaStore APIs to match Memstore APIs for backward compatibility.

Differential Revision: D80614586

fbshipit-source-id: c929fb03abf6fe21260a84df6a1b85e846da4b92

Author: aliafzal

torchrec/distributed/model_parallel.py

@@ -466,14 +466,14 @@ def get_model_tracker(self) -> ModelDeltaTracker:
         ), "Model tracker is not initialized. Add ModelTrackerConfig at DistributedModelParallel init."
         return self.model_delta_tracker
 
-    def get_delta(self, consumer: Optional[str] = None) -> Dict[str, DeltaRows]:
+    def get_unique(self, consumer: Optional[str] = None) -> Dict[str, DeltaRows]:
         """
         Returns the delta rows for the given consumer.
         """
         assert (
             self.model_delta_tracker is not None
         ), "Model tracker is not initialized. Add ModelTrackerConfig at DistributedModelParallel init."
-        return self.model_delta_tracker.get_delta(consumer)
+        return self.model_delta_tracker.get_unique(consumer)
 
     def sparse_grad_parameter_names(
         self, destination: Optional[List[str]] = None, prefix: str = ""

torchrec/distributed/model_tracker/delta_store.py

@@ -125,7 +125,7 @@ def compact(self, start_idx: int, end_idx: int) -> None:
         pass
 
     @abstractmethod
-    def get_delta(self, from_idx: int = 0) -> Dict[str, DeltaRows]:
+    def get_unique(self, from_idx: int = 0) -> Dict[str, DeltaRows]:
         """
         Return all unique/delta ids per table from the Delta Store.
 
@@ -224,7 +224,7 @@ def compact(self, start_idx: int, end_idx: int) -> None:
             )
         self.per_fqn_lookups = new_per_fqn_lookups
 
-    def get_delta(self, from_idx: int = 0) -> Dict[str, DeltaRows]:
+    def get_unique(self, from_idx: int = 0) -> Dict[str, DeltaRows]:
         r"""
         Return all unique/delta ids per table from the Delta Store.
         """

torchrec/distributed/model_tracker/model_delta_tracker.py

@@ -362,10 +362,10 @@ def get_delta_ids(self, consumer: Optional[str] = None) -> Dict[str, torch.Tenso
         Args:
             consumer (str, optional): The consumer to retrieve unique IDs for. If not specified, "default" is used as the default consumer.
         """
-        per_table_delta_rows = self.get_delta(consumer)
+        per_table_delta_rows = self.get_unique(consumer)
         return {fqn: delta_rows.ids for fqn, delta_rows in per_table_delta_rows.items()}
 
-    def get_delta(
+    def get_unique(
         self,
         consumer: Optional[str] = None,
         top_percentage: Optional[float] = 1.0,
@@ -390,7 +390,7 @@ def get_delta(
         # and index_start could be equal to index_end, in which case we should not compact again.
         if index_start < index_end:
             self.compact(index_start, index_end)
-        tracker_rows = self.store.get_delta(
+        tracker_rows = self.store.get_unique(
             from_idx=self.per_consumer_batch_idx[consumer]
         )
         self.per_consumer_batch_idx[consumer] = index_end

torchrec/distributed/model_tracker/tests/test_delta_store.py

@@ -806,8 +806,8 @@ def test_compact(self, _test_name: str, test_params: CompactTestParams) -> None:
             delta_store.compact(
                 start_idx=test_params.start_idx, end_idx=test_params.end_idx
             )
-            # Verify the result using get_delta method
-            delta_result = delta_store.get_delta()
+            # Verify the result using get_unique method
+            delta_result = delta_store.get_unique()
 
             # compare all fqns in the result
             for table_fqn, delta_rows in test_params.expected_delta.items():

torchrec/distributed/model_tracker/tests/test_model_delta_tracker.py

@@ -447,7 +447,7 @@ def test_fqn_to_feature_names(
                 ),
             ),
             (
-                "get_delta",
+                "get_unique",
                 ModelDeltaTrackerInputTestParams(
                     embedding_config_type=EmbeddingConfig,
                     embedding_tables=[
@@ -464,7 +464,7 @@ def test_fqn_to_feature_names(
                     model_tracker_config=ModelTrackerConfig(),
                 ),
                 TrackerNotInitOutputTestParams(
-                    dmp_tracker_atter="get_delta",
+                    dmp_tracker_atter="get_unique",
                 ),
             ),
         ]
@@ -1843,7 +1843,7 @@ def _test_embedding_mode(
                 tracked_out.sum().backward()
                 baseline_out.sum().backward()
 
-            delta_rows = dt.get_delta()
+            delta_rows = dt.get_unique()
 
             table_fqns = dt.fqn_to_feature_names().keys()
             table_fqns_list = list(table_fqns)
@@ -1964,7 +1964,7 @@ def _test_multiple_get(
             unittest.TestCase().assertTrue(tracked_out.allclose(baseline_out))
             tracked_out.sum().backward()
             baseline_out.sum().backward()
-            delta_rows = dt.get_delta()
+            delta_rows = dt.get_unique()
 
             # Verify that the current batch index is correct
             unittest.TestCase().assertTrue(dt.curr_batch_idx, i + 1)
@@ -2093,7 +2093,7 @@ def _test_duplication_with_momentum(
             dt_model_opt.step()
             baseline_opt.step()
 
-        delta_rows = dt.get_delta()
+        delta_rows = dt.get_unique()
         for table_fqn in table_fqns_list:
             ids = delta_rows[table_fqn].ids
             states = none_throws(delta_rows[table_fqn].states)
@@ -2162,7 +2162,7 @@ def _test_duplication_with_rowwise_adagrad(
 
         end_momentums = tbe.split_optimizer_states()[0][0].detach().clone()
 
-        delta_rows = dt.get_delta()
+        delta_rows = dt.get_unique()
         table_fqn = table_fqns_list[0]
 
         ids = delta_rows[table_fqn].ids