Renamed types to account for broader use cases (#3473)

Summary:
Pull Request resolved: #3473

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

Specific Context: Consolidating all types to a common place within torchRec

Differential Revision: D80615183

fbshipit-source-id: c730f312a1314c138f54d21b80b4ff12830f8e8f

Author: aliafzal

torchrec/distributed/model_parallel.py

@@ -30,7 +30,7 @@
 from torch.nn.parallel import DistributedDataParallel
 from torchrec.distributed.comm import get_local_size
 from torchrec.distributed.model_tracker.model_delta_tracker import ModelDeltaTrackerTrec
-from torchrec.distributed.model_tracker.types import DeltaRows, ModelTrackerConfig
+from torchrec.distributed.model_tracker.types import ModelTrackerConfig, UniqueRows
 
 from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
 from torchrec.distributed.sharding_plan import get_default_sharders
@@ -466,7 +466,7 @@ def get_model_tracker(self) -> ModelDeltaTrackerTrec:
         ), "Model tracker is not initialized. Add ModelTrackerConfig at DistributedModelParallel init."
         return self.model_delta_tracker
 
-    def get_unique(self, consumer: Optional[str] = None) -> Dict[str, DeltaRows]:
+    def get_unique(self, consumer: Optional[str] = None) -> Dict[str, UniqueRows]:
         """
         Returns the delta rows for the given consumer.
         """

torchrec/distributed/model_tracker/__init__.py

@@ -27,9 +27,9 @@
     SUPPORTED_MODULES,  # noqa
 )
 from torchrec.distributed.model_tracker.types import (
-    DeltaRows,  # noqa
-    EmbdUpdateMode,  # noqa
     IndexedLookup,  # noqa
     ModelTrackerConfig,  # noqa
     TrackingMode,  # noqa
+    UniqueRows,  # noqa
+    UpdateMode,  # noqa
 )

torchrec/distributed/model_tracker/delta_store.py

@@ -13,34 +13,34 @@
 
 import torch
 from torchrec.distributed.model_tracker.types import (
-    DeltaRows,
-    EmbdUpdateMode,
     IndexedLookup,
+    UniqueRows,
+    UpdateMode,
 )
 from torchrec.distributed.utils import none_throws
 
 
 def _compute_unique_rows(
     ids: List[torch.Tensor],
     states: Optional[List[torch.Tensor]],
-    mode: EmbdUpdateMode,
-) -> DeltaRows:
+    mode: UpdateMode,
+) -> UniqueRows:
     r"""
     To calculate unique ids and embeddings
     """
-    if mode == EmbdUpdateMode.NONE:
-        assert states is None, f"{mode=} == EmbdUpdateMode.NONE but received embeddings"
+    if mode == UpdateMode.NONE:
+        assert states is None, f"{mode=} == UpdateMode.NONE but received embeddings"
         unique_ids = torch.cat(ids).unique(return_inverse=False)
-        return DeltaRows(ids=unique_ids, states=None)
+        return UniqueRows(ids=unique_ids, states=None)
     else:
         assert (
             states is not None
-        ), f"{mode=} != EmbdUpdateMode.NONE but received no embeddings"
+        ), f"{mode=} != UpdateMode.NONE but received no embeddings"
 
         cat_ids = torch.cat(ids)
         cat_states = torch.cat(states)
 
-        if mode == EmbdUpdateMode.LAST:
+        if mode == UpdateMode.LAST:
             cat_ids = cat_ids.flip(dims=[0])
             cat_states = cat_states.flip(dims=[0])
 
@@ -65,7 +65,7 @@ def _compute_unique_rows(
 
         # Use first occurrence indices to select corresponding embedding row.
         unique_states = cat_states[first_occurrence]
-        return DeltaRows(ids=unique_ids, states=unique_states)
+        return UniqueRows(ids=unique_ids, states=unique_states)
 
 
 class DeltaStore(ABC):
@@ -81,7 +81,7 @@ class DeltaStore(ABC):
     """
 
     @abstractmethod
-    def __init__(self, embdUpdateMode: EmbdUpdateMode = EmbdUpdateMode.NONE) -> None:
+    def __init__(self, updateMode: UpdateMode = UpdateMode.NONE) -> None:
         pass
 
     @abstractmethod
@@ -125,7 +125,7 @@ def compact(self, start_idx: int, end_idx: int) -> None:
         pass
 
     @abstractmethod
-    def get_unique(self, from_idx: int = 0) -> Dict[str, DeltaRows]:
+    def get_unique(self, from_idx: int = 0) -> Dict[str, UniqueRows]:
         """
         Return all unique/delta ids per table from the Delta Store.
 
@@ -151,9 +151,9 @@ class DeltaStoreTrec(DeltaStore):
     how to handle duplicate ids when compacting or retrieving embeddings.
     """
 
-    def __init__(self, embdUpdateMode: EmbdUpdateMode = EmbdUpdateMode.NONE) -> None:
-        super().__init__(embdUpdateMode)
-        self.embdUpdateMode = embdUpdateMode
+    def __init__(self, updateMode: UpdateMode = UpdateMode.NONE) -> None:
+        super().__init__(updateMode)
+        self.updateMode = updateMode
         self.per_fqn_lookups: Dict[str, List[IndexedLookup]] = {}
 
     def append(
@@ -205,11 +205,11 @@ def compact(self, start_idx: int, end_idx: int) -> None:
             ids = [lookup.ids for lookup in lookups_to_compact]
             states = (
                 [none_throws(lookup.states) for lookup in lookups_to_compact]
-                if self.embdUpdateMode != EmbdUpdateMode.NONE
+                if self.updateMode != UpdateMode.NONE
                 else None
             )
             delta_rows = _compute_unique_rows(
-                ids=ids, states=states, mode=self.embdUpdateMode
+                ids=ids, states=states, mode=self.updateMode
             )
             new_per_fqn_lookups[table_fqn] = (
                 lookups[:index_l]
@@ -224,12 +224,12 @@ def compact(self, start_idx: int, end_idx: int) -> None:
             )
         self.per_fqn_lookups = new_per_fqn_lookups
 
-    def get_unique(self, from_idx: int = 0) -> Dict[str, DeltaRows]:
+    def get_unique(self, from_idx: int = 0) -> Dict[str, UniqueRows]:
         r"""
         Return all unique/delta ids per table from the Delta Store.
         """
 
-        delta_per_table_fqn: Dict[str, DeltaRows] = {}
+        delta_per_table_fqn: Dict[str, UniqueRows] = {}
         for table_fqn, lookups in self.per_fqn_lookups.items():
             compact_ids = [
                 lookup.ids for lookup in lookups if lookup.batch_idx >= from_idx
@@ -240,11 +240,11 @@ def get_unique(self, from_idx: int = 0) -> Dict[str, DeltaRows]:
                     for lookup in lookups
                     if lookup.batch_idx >= from_idx
                 ]
-                if self.embdUpdateMode != EmbdUpdateMode.NONE
+                if self.updateMode != UpdateMode.NONE
                 else None
             )
 
             delta_per_table_fqn[table_fqn] = _compute_unique_rows(
-                ids=compact_ids, states=compact_states, mode=self.embdUpdateMode
+                ids=compact_ids, states=compact_states, mode=self.updateMode
             )
         return delta_per_table_fqn

torchrec/distributed/model_tracker/model_delta_tracker.py

@@ -29,36 +29,36 @@
 from torchrec.distributed.embeddingbag import ShardedEmbeddingBagCollection
 from torchrec.distributed.model_tracker.delta_store import DeltaStoreTrec
 from torchrec.distributed.model_tracker.types import (
-    DeltaRows,
-    EmbdUpdateMode,
     TrackingMode,
+    UniqueRows,
+    UpdateMode,
 )
 from torchrec.distributed.utils import none_throws
 
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 
-UPDATE_MODE_MAP: Dict[TrackingMode, EmbdUpdateMode] = {
+UPDATE_MODE_MAP: Dict[TrackingMode, UpdateMode] = {
     # Only IDs are tracked, no additional state is stored.
-    TrackingMode.ID_ONLY: EmbdUpdateMode.NONE,
-    # TrackingMode.EMBEDDING utilizes EmbdUpdateMode.FIRST to ensure that
+    TrackingMode.ID_ONLY: UpdateMode.NONE,
+    # TrackingMode.EMBEDDING utilizes UpdateMode.FIRST to ensure that
     # the earliest embedding values are stored since the last checkpoint
     # or snapshot. This mode is used for computing topk delta rows, which
     # is currently achieved by running (new_emb - old_emb).norm().topk().
-    TrackingMode.EMBEDDING: EmbdUpdateMode.FIRST,
-    # TrackingMode.MOMENTUM utilizes EmbdUpdateMode.LAST to ensure that
+    TrackingMode.EMBEDDING: UpdateMode.FIRST,
+    # TrackingMode.MOMENTUM utilizes UpdateMode.LAST to ensure that
     # the most recent momentum values—capturing the accumulated gradient
     # direction and magnitude—are stored since the last batch.
     # This mode supports approximate top-k delta-row selection, can be
     # obtained by running momentum.norm().topk().
-    TrackingMode.MOMENTUM_LAST: EmbdUpdateMode.LAST,
+    TrackingMode.MOMENTUM_LAST: UpdateMode.LAST,
     # MOMENTUM_DIFF keeps a running sum of the square of the gradients per row.
     # Within each publishing interval, we track the starting value of this running
     # sum on all used rows and then do a lookup when ``get_delta`` is called to query
     # the latest sum. Then we can compute the delta of the two values and return them
     # together with the row ids.
-    TrackingMode.MOMENTUM_DIFF: EmbdUpdateMode.FIRST,
+    TrackingMode.MOMENTUM_DIFF: UpdateMode.FIRST,
     # The same as MOMENTUM_DIFF. Adding for backward compatibility.
-    TrackingMode.ROWWISE_ADAGRAD: EmbdUpdateMode.FIRST,
+    TrackingMode.ROWWISE_ADAGRAD: UpdateMode.FIRST,
 }
 
 # Tracking is current only supported for ShardedEmbeddingCollection and ShardedEmbeddingBagCollection.
@@ -111,7 +111,7 @@ def get_unique(
         top_percentage: Optional[float] = 1.0,
         per_table_percentage: Optional[Dict[str, Tuple[float, str]]] = None,
         sorted_by_indices: Optional[bool] = True,
-    ) -> Dict[str, DeltaRows]:
+    ) -> Dict[str, UniqueRows]:
         """
         Return a dictionary of hit local IDs and parameter states / embeddings for each sparse feature.
 
@@ -439,7 +439,7 @@ def get_unique(
         top_percentage: Optional[float] = 1.0,
         per_table_percentage: Optional[Dict[str, Tuple[float, str]]] = None,
         sorted_by_indices: Optional[bool] = True,
-    ) -> Dict[str, DeltaRows]:
+    ) -> Dict[str, UniqueRows]:
         """
         Return a dictionary of hit local IDs and parameter states / embeddings for each sparse feature. The Values are first keyed by submodule FQN.
 
@@ -471,8 +471,7 @@ def get_unique(
                 assert (
                     fqn in square_sum_map
                 ), f"{fqn} not found in {square_sum_map.keys()}"
-                # pyre-fixme[58]: `-` is not supported for operand types `Tensor`
-                #  and `Optional[Tensor]`.
+                # pyre-fixme[58]: `-` is not supported for operand types `Tensor` and `Optional[Tensor]`.
                 rows.states = square_sum_map[fqn][rows.ids] - rows.states
 
         return tracker_rows