Add raw_id_tracker for tracking hash_zch_identities of MPZCH module (#3501)

Summary:

This diff introduces a new `RawIdTracker` class that extends TorchRec's model delta tracking infra to capture and track raw hash identities from MCC  modules during training. This is specifically required for tracking raw ids for MPZCH tables.

internal
This is needed to support MPZCH modules for Raw embedding streaming.
Mode details :
https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld

Differential Revision: D84920167

Author: aliafzal

torchrec/distributed/embedding.py

@@ -1605,7 +1605,7 @@ def compute_and_output_dist(
             ):
                 embs = lookup(features)
                 if self.post_lookup_tracker_fn is not None:
-                    self.post_lookup_tracker_fn(features, embs, self)
+                    self.post_lookup_tracker_fn(features, embs, self, None)
 
             with maybe_annotate_embedding_event(
                 EmbeddingEvent.OUTPUT_DIST, self._module_fqn, sharding_type

torchrec/distributed/embedding_types.py

@@ -391,7 +391,15 @@ def __init__(
         self._lookups: List[nn.Module] = []
         self._output_dists: List[nn.Module] = []
         self.post_lookup_tracker_fn: Optional[
-            Callable[[KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None]
+            Callable[
+                [
+                    KeyedJaggedTensor,
+                    torch.Tensor,
+                    Optional[nn.Module],
+                    Optional[torch.Tensor],
+                ],
+                None,
+            ]
         ] = None
         self.post_odist_tracker_fn: Optional[Callable[..., None]] = None
 
@@ -445,7 +453,13 @@ def train(self, mode: bool = True):  # pyre-ignore[3]
     def register_post_lookup_tracker_fn(
         self,
         record_fn: Callable[
-            [KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None
+            [
+                KeyedJaggedTensor,
+                torch.Tensor,
+                Optional[nn.Module],
+                Optional[torch.Tensor],
+            ],
+            None,
         ],
     ) -> None:
         """

torchrec/distributed/embeddingbag.py

@@ -1671,7 +1671,7 @@ def compute_and_output_dist(
             ):
                 embs = lookup(features)
                 if self.post_lookup_tracker_fn is not None:
-                    self.post_lookup_tracker_fn(features, embs, self)
+                    self.post_lookup_tracker_fn(features, embs, self, None)
 
             with maybe_annotate_embedding_event(
                 EmbeddingEvent.OUTPUT_DIST,

torchrec/distributed/mc_modules.py

@@ -238,7 +238,15 @@ def __init__(
         self._use_index_dedup = use_index_dedup
         self._initialize_torch_state()
         self.post_lookup_tracker_fn: Optional[
-            Callable[[KeyedJaggedTensor, torch.Tensor], None]
+            Callable[
+                [
+                    KeyedJaggedTensor,
+                    torch.Tensor,
+                    Optional[nn.Module],
+                    Optional[torch.Tensor],
+                ],
+                None,
+            ]
         ] = None
 
     def _initialize_torch_state(self) -> None:
@@ -756,6 +764,8 @@ def compute(
                         if self.post_lookup_tracker_fn is not None:
                             self.post_lookup_tracker_fn(
                                 KeyedJaggedTensor.from_jt_dict(output),
+                                torch.empty(0),
+                                None,
                                 mcm._hash_zch_identities.index_select(
                                     dim=0, index=mc_input[table].values()
                                 ),
@@ -782,6 +792,8 @@ def compute(
                     if self.post_lookup_tracker_fn is not None:
                         self.post_lookup_tracker_fn(
                             KeyedJaggedTensor.from_jt_dict(mc_input),
+                            torch.empty(0),
+                            None,
                             mcm._hash_zch_identities.index_select(dim=0, index=values),
                         )
 
@@ -876,7 +888,15 @@ def unsharded_module_type(self) -> Type[ManagedCollisionCollection]:
 
     def register_post_lookup_tracker_fn(
         self,
-        record_fn: Callable[[KeyedJaggedTensor, torch.Tensor], None],
+        record_fn: Callable[
+            [
+                KeyedJaggedTensor,
+                torch.Tensor,
+                Optional[nn.Module],
+                Optional[torch.Tensor],
+            ],
+            None,
+        ],
     ) -> None:
         """
         Register a function to be called after lookup is done. This is used for

torchrec/distributed/model_tracker/delta_store.py

@@ -91,6 +91,7 @@ def append(
         fqn: str,
         ids: torch.Tensor,
         states: Optional[torch.Tensor],
+        raw_ids: Optional[torch.Tensor] = None,
     ) -> None:
         """
         Append a batch of ids and states to the store for a specific table.
@@ -162,10 +163,11 @@ def append(
         fqn: str,
         ids: torch.Tensor,
         states: Optional[torch.Tensor],
+        raw_ids: Optional[torch.Tensor] = None,
     ) -> None:
         table_fqn_lookup = self.per_fqn_lookups.get(fqn, [])
         table_fqn_lookup.append(
-            IndexedLookup(batch_idx=batch_idx, ids=ids, states=states)
+            IndexedLookup(batch_idx=batch_idx, ids=ids, states=states, raw_ids=raw_ids)
         )
         self.per_fqn_lookups[fqn] = table_fqn_lookup
 
@@ -224,6 +226,20 @@ def compact(self, start_idx: int, end_idx: int) -> None:
             )
         self.per_fqn_lookups = new_per_fqn_lookups
 
+    def get_indexed_lookups(
+        self, start_idx: int, end_idx: int
+    ) -> Dict[str, List[IndexedLookup]]:
+        r"""
+        Return all unique/delta ids per table from the Delta Store.
+        """
+        per_fqn_lookups: Dict[str, List[IndexedLookup]] = {}
+        for table_fqn, lookups in self.per_fqn_lookups.items():
+            indexices = [h.batch_idx for h in lookups]
+            index_l = bisect_left(indexices, start_idx)
+            index_r = bisect_left(indexices, end_idx)
+            per_fqn_lookups[table_fqn] = lookups[index_l:index_r]
+        return per_fqn_lookups
+
     def get_unique(self, from_idx: int = 0) -> Dict[str, UniqueRows]:
         r"""
         Return all unique/delta ids per table from the Delta Store.

torchrec/distributed/model_tracker/model_delta_tracker.py

@@ -84,6 +84,7 @@ def record_lookup(
         kjt: KeyedJaggedTensor,
         states: torch.Tensor,
         emb_module: Optional[nn.Module] = None,
+        raw_ids: Optional[torch.Tensor] = None,
     ) -> None:
         """
         Records the IDs from a given KeyedJaggedTensor and their corresponding embeddings/parameter states.
@@ -131,6 +132,13 @@ def clear(self, consumer: Optional[str] = None) -> None:
         """
         pass
 
+    @abstractmethod
+    def step(self) -> None:
+        """
+        Advance the batch index for all consumers.
+        """
+        pass
+
 
 class ModelDeltaTrackerTrec(ModelDeltaTracker):
     r"""
@@ -244,6 +252,7 @@ def record_lookup(
         kjt: KeyedJaggedTensor,
         states: torch.Tensor,
         emb_module: Optional[nn.Module] = None,
+        raw_ids: Optional[torch.Tensor] = None,
     ) -> None:
         """
         Records the IDs from a given KeyedJaggedTensor and their corresponding embeddings/parameter states.