Delta tracker DMP integration (#3064)

Summary:

## This Diff
Adds ModelDeltaTracker integration with DMP (DistributedModelParallel) and sharded modules. This integration enables tracking of embedding IDs, embeddings, and optimizer states during model execution, which is particularly useful for online training scenarios.

### Key Components:
**ModelTrackerConfig Integration**:
* Added ModelTrackerConfig parameter to DMP constructor
* When provided, automatically initializes ModelDeltaTracker
*  Configurable options include tracking_mode, delete_on_read, auto_compact, and fqns_to_skip

**Custom Callables for Tracking**:
* Added custom post_lookup_hook in ShardedModule to capture IDs and embeddings after lookup operations. This provides tracking ids/states natively into torchrec without registering any nn.Module specific hooks.
* Added post_odist_hook for auto-compaction of tracked data. This custom hook provides native support for overlapping compaction with odist.
* Implemented pre_forward callables in DMP for operations like batch index incrementation

**Model Parallel API Enhancements**:
* Added `get_model_tracker()` method to DistributedModelParallel for direct access to the ModelDeltaTracker instance. This API give the flexibility to integrate model tracker into required components directly without needing to access the dmp_module.
* Added `get_delta()` method as a convenience API to retrieve delta rows from dmp_module. 

**Embedding Module Changes**:
* Enhanced ShardedEmbeddingBag and ShardedEmbedding to support tracking hooks / Callable
* Added hook registration methods in embedding modules
* Implemented tracking support for different optimizer states (momentum, Adam states)
 

## ModelDeltaTracker Context
ModelDeltaTracker is a utility for tracking and retrieving unique IDs and their corresponding embeddings or states from embedding modules in model using Torchrec. It's particularly useful for:

1. Identifying which embedding rows were accessed during model execution
2. Retrieving the latest delta or unique rows for a model
3. Computing top-k changed embeddings
4. Supporting streaming updated embeddings between systems during online training


For more details see diff:D75853147 or PR #3057

Differential Revision: D76202371

Author: aliafzal

torchrec/distributed/embedding.py

@@ -1514,13 +1514,17 @@ def compute_and_output_dist(
                 EmbeddingEvent.LOOKUP, self._module_fqn, sharding_type
             ):
                 embs = lookup(features)
+                if self.post_lookup_tracker_hook is not None:
+                    self.post_lookup_tracker_hook(self, features, embs)
 
             with maybe_annotate_embedding_event(
                 EmbeddingEvent.OUTPUT_DIST, self._module_fqn, sharding_type
             ):
                 awaitables_per_sharding.append(
                     odist(embs.view(-1, embedding_dim), sharding_ctx)
                 )
+                if self.post_odist_tracker_hook is not None:
+                    self.post_odist_tracker_hook()
 
             features_before_all2all_per_sharding.append(
                 # pyre-fixme[6]: For 1st argument expected `KeyedJaggedTensor` but

torchrec/distributed/embedding_types.py

@@ -9,10 +9,12 @@
 
 import abc
 import copy
+import logging as logger
 from dataclasses import dataclass
 from enum import Enum, unique
 from typing import (
     Any,
+    Callable,
     Dict,
     Generic,
     Iterable,
@@ -370,6 +372,10 @@ def __init__(
         self._input_dists: List[nn.Module] = []
         self._lookups: List[nn.Module] = []
         self._output_dists: List[nn.Module] = []
+        self.post_lookup_tracker_hook: Optional[
+            Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]
+        ] = None
+        self.post_odist_tracker_hook: Optional[Callable[..., None]] = None
 
     def prefetch(
         self,
@@ -418,6 +424,41 @@ def train(self, mode: bool = True):  # pyre-ignore[3]
 
         return self
 
+    def register_post_lookup_tracker_hook(
+        self,
+        record_fn: Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None],
+    ) -> None:
+        """
+        Register a hook to be called after lookup is done. This is used for
+        tracking the lookup results and optimizer states.
+
+        Args:
+            record_fn (Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]): A custom record function to be called after lookup is done.
+
+        """
+        if self.post_lookup_tracker_hook is not None:
+            logger.warning(
+                "[ModelDeltaTracker] Custom record function already defined, overriding with new callable"
+            )
+        self.post_lookup_tracker_hook = record_fn
+
+    def register_post_odist_tracker_hook(
+        self,
+        record_fn: Callable[..., None],
+    ) -> None:
+        """
+        Register a hook to be called after registering odist awaitable.
+
+        Args:
+            record_fn (Callable[Callable[..., None]):
+
+        """
+        if self.post_odist_tracker_hook is not None:
+            logger.warning(
+                "[ModelDeltaTracker] Compaction function already defined, overriding with new callable"
+            )
+        self.post_odist_tracker_hook = record_fn
+
     @property
     def unsharded_module_type(self) -> Type[nn.Module]:
         """

torchrec/distributed/embeddingbag.py

@@ -1458,13 +1458,17 @@ def compute_and_output_dist(
                 sharding_type,
             ):
                 embs = lookup(features)
+                if self.post_lookup_tracker_hook is not None:
+                    self.post_lookup_tracker_hook(self, features, embs)
 
             with maybe_annotate_embedding_event(
                 EmbeddingEvent.OUTPUT_DIST,
                 self._module_fqn,
                 sharding_type,
             ):
                 awaitables.append(dist(embs, sharding_context))
+                if self.post_odist_tracker_hook is not None:
+                    self.post_odist_tracker_hook()
 
             if sharding_context:
                 batch_size_per_feature_pre_a2a.extend(

torchrec/distributed/model_parallel.py

@@ -29,6 +29,11 @@
 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,
+    SUPPORTED_MODULES,
+)
+from torchrec.distributed.model_tracker.types import DeltaRows, ModelTrackerConfig
 
 from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
 from torchrec.distributed.sharding_plan import get_default_sharders
@@ -208,6 +213,7 @@ class DistributedModelParallel(nn.Module, FusedOptimizerModule):
         init_parameters (bool): initialize parameters for modules still on meta device.
         data_parallel_wrapper (Optional[DataParallelWrapper]): custom wrapper for data
             parallel modules.
+        model_tracker_config (Optional[DeltaTrackerConfig]): config for model tracker.
 
     Example::
 
@@ -234,6 +240,7 @@ def __init__(
         init_data_parallel: bool = True,
         init_parameters: bool = True,
         data_parallel_wrapper: Optional[DataParallelWrapper] = None,
+        model_tracker_config: Optional[ModelTrackerConfig] = None,
     ) -> None:
         super().__init__()
         torch._C._log_api_usage_once(f"torchrec.distributed.{self.__class__.__name__}")
@@ -242,6 +249,11 @@ def __init__(
 
         self._ddp_wrapped: bool = False
 
+        self.has_model_tracker: bool = model_tracker_config is not None
+
+        # List of callables to be executed before forward
+        self._pre_forward_callables: List[Callable[..., None]] = []
+
         if env is None:
             pg = dist.GroupMember.WORLD
             assert pg is not None, "Process group is not initialized"
@@ -286,6 +298,26 @@ def __init__(
         if init_data_parallel:
             self.init_data_parallel()
 
+        if model_tracker_config is not None:
+            self.model_delta_tracker: ModelDeltaTracker = self._init_delta_tracker(
+                model_tracker_config, self._dmp_wrapped_module
+            )
+            tracked_modules = self.model_delta_tracker.get_tracked_modules()
+            for module in tracked_modules.values():
+                if isinstance(module, SUPPORTED_MODULES):
+                    # register post lookup hook
+                    module.register_post_lookup_tracker_hook(
+                        self.model_delta_tracker.record_lookup
+                    )
+                    # register auto compaction hook at odist
+                    if model_tracker_config.auto_compact:
+                        module.register_post_odist_tracker_hook(
+                            self.model_delta_tracker.trigger_compaction
+                        )
+            self._pre_forward_callables.append(
+                self.model_delta_tracker.increment_batch_idx
+            )
+
     @property
     def module(self) -> nn.Module:
         """
@@ -307,8 +339,25 @@ def module(self, value: nn.Module) -> None:
 
     # pyre-ignore [2, 3]
     def forward(self, *args, **kwargs) -> Any:
+        # Execute pre-forward callables. Currently used to update batch index
+        # for model tracker.
+        for callable_fn in self._pre_forward_callables:
+            callable_fn()
+
         return self._dmp_wrapped_module(*args, **kwargs)
 
+    def register_pre_forward_callable(self, callable_fn: Callable[..., None]) -> None:
+        """
+        Register a callable to be executed before the forward pass.
+
+        The callable can accept kwargs that match the forward method's kwargs,
+        or it can be called without any arguments.
+
+        Args:
+            callable_fn: A callable function to be executed before forward
+        """
+        self._pre_forward_callables.append(callable_fn)
+
     def init_data_parallel(self) -> None:
         """
         See init_data_parallel c-tor argument for usage.
@@ -344,6 +393,19 @@ def copy(
     def _init_dmp(self, module: nn.Module) -> nn.Module:
         return self._shard_modules_impl(module)
 
+    def _init_delta_tracker(
+        self, model_tracker_config: ModelTrackerConfig, module: nn.Module
+    ) -> ModelDeltaTracker:
+        # Init delta tracker if config is provided
+        return ModelDeltaTracker(
+            model=module,
+            consumers=model_tracker_config.consumers,
+            delete_on_read=model_tracker_config.delete_on_read,
+            auto_compact=model_tracker_config.auto_compact,
+            mode=model_tracker_config.tracking_mode,
+            fqns_to_skip=model_tracker_config.fqns_to_skip,
+        )
+
     def _init_optim(self, module: nn.Module) -> CombinedOptimizer:
         # pyre-ignore [6]
         return CombinedOptimizer(self._fused_optim_impl(module, []))
@@ -421,6 +483,25 @@ def init_parameters(module: nn.Module) -> None:
 
         module.apply(init_parameters)
 
+    def get_model_tracker(self) -> ModelDeltaTracker:
+        """
+        Returns the model tracker if it exists.
+        """
+
+        assert (
+            self.has_model_tracker
+        ), "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]:
+        """
+        Returns the delta rows for the given consumer.
+        """
+        assert (
+            self.has_model_tracker
+        ), "Model tracker is not initialized. Add ModelTrackerConfig at DistributedModelParallel init."
+        return self.model_delta_tracker.get_delta(consumer)
+
     def sparse_grad_parameter_names(
         self, destination: Optional[List[str]] = None, prefix: str = ""
     ) -> List[str]:

torchrec/distributed/model_tracker/model_delta_tracker.py

@@ -104,6 +104,37 @@ def __init__(
                 self.feature_to_fqn[feature_name] = fqn
         logger.info(f"feature_to_fqn: {self.feature_to_fqn}")
 
+    def increment_batch_idx(self) -> None:
+        self.curr_batch_idx += 1
+
+    def trigger_compaction(self) -> None:
+        if self.curr_compact_index >= self.curr_batch_idx:
+            # only trigger compaction once per iteration
+            return
+
+        self.curr_compact_index += 1
+        # TODO: May need to revisit the compaction logic with multiple consmers.
+        # At present we take the max per_consumer_batch_idx to ensure we only compact
+        # newely received lookups
+
+        # The trigger_compaction() function is expected to overlap with comms to hide
+        # compaction compute overhead. Currently, we overlap compaction with odist
+        # because ID tracking occurs during local embedding lookup, which takes place
+        # before odist. This way, auto_compact always merges all past IDs tensors since
+        # the last get_delta call into a single IDs tensor per FQN.
+        #
+        # For delete_on_read=True, get_delta() should delete up to per_consumer_batch_idx
+        # (exclusive). So the compaction should start from per_consumer_batch_idx.
+        #
+        # For delete_on_read=False, get_delta() won't delete tensors, but it does advance
+        # per_consumer_batch_idx accordingly, where all ids prior to per_consumer_batch_idx (exclusive)
+        # should have been compacted into one tensor regardless of auto_compact=True/False.
+        # Therefore, all future compactions should start from per_consumer_batch_idx.
+        start_idx = max(self.per_consumer_batch_idx.values())
+        end_idx = self.curr_batch_idx
+        if start_idx < end_idx:
+            self.compact(start_idx=start_idx, end_idx=end_idx)
+
     def record_lookup(
         self, emb_module: nn.Module, kjt: KeyedJaggedTensor, states: torch.Tensor
     ) -> None:
@@ -329,8 +360,13 @@ def compact(self, start_idx: int, end_idx: int) -> None:
         self.store.compact(start_idx, end_idx)
 
     def _clean_fqn_fn(self, fqn: str) -> str:
-        # strip DMP internal module FQN prefix to match state dict FQN
-        return fqn.replace("_dmp_wrapped_module.module.", "")
+        # strip internal module FQN prefixes to match state dict FQN
+        # handles both "_dmp_wrapped_module.module." and "module." prefixes
+        prefixes_to_strip = ["_dmp_wrapped_module.module.", "module."]
+        for prefix in prefixes_to_strip:
+            if fqn.startswith(prefix):
+                return fqn[len(prefix) :]
+        return fqn
 
     def _validate_mode(self) -> None:
         "To validate the mode is supported for the given module"