Update ModelDeltaTracker to be Generic (#3470)

Summary:
Pull Request resolved: #3470

Make ModelDeltaTracker generic to allow use case specific custom implementations

internal

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

Specific Context: DeltaTracker implements primitives to allow tracking of embedding ids and states to optimize checkpointing and embedding freshness. As part of transitioning to a common DeltaTracker, we are adding a generic ModelDeltaTracker. MRS DeltaTracker will extend from Generic ModelDeltaTracker.

Differential Revision: D80614689

fbshipit-source-id: 07126b50f4c2819f7492ee0fdb549533f102e970

Author: aliafzal

torchrec/distributed/model_parallel.py

@@ -29,7 +29,7 @@
 from torch.nn.modules.module import _IncompatibleKeys
 from torch.nn.parallel import DistributedDataParallel
 from torchrec.distributed.comm import get_local_size
-from torchrec.distributed.model_tracker.model_delta_tracker import ModelDeltaTracker
+from torchrec.distributed.model_tracker.model_delta_tracker import ModelDeltaTrackerTrec
 from torchrec.distributed.model_tracker.types import DeltaRows, ModelTrackerConfig
 
 from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
@@ -293,7 +293,7 @@ def __init__(
         if init_data_parallel:
             self.init_data_parallel()
 
-        self.model_delta_tracker: Optional[ModelDeltaTracker] = (
+        self.model_delta_tracker: Optional[ModelDeltaTrackerTrec] = (
             self._init_delta_tracker(model_tracker_config, self._dmp_wrapped_module)
             if model_tracker_config is not None
             else None
@@ -369,9 +369,9 @@ def _init_dmp(self, module: nn.Module) -> nn.Module:
 
     def _init_delta_tracker(
         self, model_tracker_config: ModelTrackerConfig, module: nn.Module
-    ) -> ModelDeltaTracker:
+    ) -> ModelDeltaTrackerTrec:
         # Init delta tracker if config is provided
-        return ModelDeltaTracker(
+        return ModelDeltaTrackerTrec(
             model=module,
             consumers=model_tracker_config.consumers,
             delete_on_read=model_tracker_config.delete_on_read,
@@ -456,7 +456,7 @@ def init_parameters(module: nn.Module) -> None:
 
         module.apply(init_parameters)
 
-    def get_model_tracker(self) -> ModelDeltaTracker:
+    def get_model_tracker(self) -> ModelDeltaTrackerTrec:
         """
         Returns the model tracker if it exists.
         """

torchrec/distributed/model_tracker/model_delta_tracker.py

@@ -7,6 +7,7 @@
 
 # pyre-strict
 import logging as logger
+from abc import ABC, abstractmethod
 from collections import Counter, OrderedDict
 from typing import Dict, Iterable, List, Optional, Tuple
 
@@ -64,10 +65,73 @@
 SUPPORTED_MODULES = (ShardedEmbeddingCollection, ShardedEmbeddingBagCollection)
 
 
-class ModelDeltaTracker:
+class ModelDeltaTracker(ABC):
     r"""
 
-    ModelDeltaTracker provides a way to track and retrieve unique IDs for supported modules, along with optional support
+    Abstract base class for ModelDeltaTracker that provides a way to track and retrieve unique IDs for supported modules,
+    along with optional support for tracking corresponding embeddings or states. This is useful for identifying and
+    retrieving the latest delta or unique rows for a given model, which can help compute topk or to stream updated
+    embeddings from predictors to trainers during online training.
+
+    """
+
+    DEFAULT_CONSUMER: str = "default"
+
+    @abstractmethod
+    def record_lookup(
+        self, emb_module: nn.Module, kjt: KeyedJaggedTensor, states: torch.Tensor
+    ) -> None:
+        """
+        Records the IDs from a given KeyedJaggedTensor and their corresponding embeddings/parameter states.
+
+        Args:
+            emb_module (nn.Module): The embedding module in which the lookup was performed.
+            kjt (KeyedJaggedTensor): The KeyedJaggedTensor containing IDs to record.
+            states (torch.Tensor): The embeddings or states corresponding to the IDs in the kjt.
+        """
+        pass
+
+    @abstractmethod
+    def get_unique_ids(self, consumer: Optional[str] = None) -> Dict[str, torch.Tensor]:
+        """
+        Return a dictionary of hit local IDs for each sparse feature.
+
+        Args:
+            consumer (str, optional): The consumer to retrieve unique IDs for.
+        """
+        pass
+
+    @abstractmethod
+    def get_unique(
+        self,
+        consumer: Optional[str] = None,
+        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]:
+        """
+        Return a dictionary of hit local IDs and parameter states / embeddings for each sparse feature.
+
+        Args:
+            consumer (str, optional): The consumer to retrieve delta values for.
+        """
+        pass
+
+    @abstractmethod
+    def clear(self, consumer: Optional[str] = None) -> None:
+        """
+        Clear tracked IDs for a given consumer.
+
+        Args:
+            consumer (str, optional): The consumer to clear IDs/States for.
+        """
+        pass
+
+
+class ModelDeltaTrackerTrec(ModelDeltaTracker):
+    r"""
+
+    ModelDeltaTrackerTrec provides a way to track and retrieve unique IDs for supported modules, along with optional support
     for tracking corresponding embeddings or states. This is useful for identifying and retrieving the latest delta or
     unique rows for a given model, which can help compute topk or to stream updated embeddings from predictors to trainers during
     online training. Unique IDs or states can be retrieved by calling the get_delta() method.
@@ -85,8 +149,6 @@ class ModelDeltaTracker:
 
     """
 
-    DEFAULT_CONSUMER: str = "default"
-
     def __init__(
         self,
         model: nn.Module,
@@ -354,7 +416,7 @@ def get_latest(self) -> Dict[str, torch.Tensor]:
 
         return ret
 
-    def get_delta_ids(self, consumer: Optional[str] = None) -> Dict[str, torch.Tensor]:
+    def get_unique_ids(self, consumer: Optional[str] = None) -> Dict[str, torch.Tensor]:
         """
         Return a dictionary of hit local IDs for each sparse feature. Ids are
         first keyed by submodule FQN.

torchrec/distributed/model_tracker/tests/test_model_delta_tracker.py

@@ -1739,7 +1739,7 @@ def _test_id_mode(
             tracked_out.sum().backward()
             baseline_out.sum().backward()
 
-        delta_ids = dt.get_delta_ids()
+        delta_ids = dt.get_unique_ids()
 
         table_fqns = dt.fqn_to_feature_names().keys()
 
@@ -2035,7 +2035,7 @@ def _test_multiple_consumer(
             unittest.TestCase().assertTrue(tracked_out.allclose(baseline_out))
             tracked_out.sum().backward()
             baseline_out.sum().backward()
-            delta_rows = dt.get_delta_ids(consumer=consumer)
+            delta_rows = dt.get_unique_ids(consumer=consumer)
 
             # Verify that the current batch index is correct
             unittest.TestCase().assertTrue(dt.curr_batch_idx, i + 1)