Store and retrieve hash_zch_runtime_meta within mc_embbedding submodules

Summary:
See https://fb.workplace.com/groups/1404957374198553/permalink/1610214197006202/

Similar to D85999577, we store the `hash_zch_runtime_meta` when it is being look up in `mc_modules.py` in `raw_id_tracker` then access it in `batched_embedding_kernel.py` which will then be streamed to the inference side (see the other diffs in this stack D87810125)

Differential Revision: D88623165

Author: Joey Yang

torchrec/distributed/batched_embedding_kernel.py

@@ -1764,19 +1764,29 @@ def init_parameters(self) -> None:
                     weight_init_max,
                 )
 
-    def forward(self, features: KeyedJaggedTensor) -> torch.Tensor:
-        hash_zch_identities = self._get_hash_zch_identities(features)
-        if hash_zch_identities is None:
+    def forward(
+        self,
+        features: KeyedJaggedTensor,
+    ) -> torch.Tensor:
+        forward_args: Dict[str, Any] = {}
+        identities_and_metadata = self._get_hash_zch_identities_and_metadata(features)
+        if identities_and_metadata is not None:
+            hash_zch_identities, hash_zch_runtime_meta = identities_and_metadata
+            forward_args["hash_zch_identities"] = hash_zch_identities
+            if hash_zch_runtime_meta is not None:
+                forward_args["hash_zch_runtime_meta"] = hash_zch_runtime_meta
+
+        if len(forward_args) == 0:
             return self.emb_module(
                 indices=features.values().long(),
                 offsets=features.offsets().long(),
             )
-
-        return self.emb_module(
-            indices=features.values().long(),
-            offsets=features.offsets().long(),
-            hash_zch_identities=hash_zch_identities,
-        )
+        else:
+            return self.emb_module(
+                indices=features.values().long(),
+                offsets=features.offsets().long(),
+                **forward_args,
+            )
 
     # pyre-fixme[14]: `state_dict` overrides method defined in `Module` inconsistently.
     def state_dict(
@@ -2841,9 +2851,12 @@ def forward(
         features: KeyedJaggedTensor,
     ) -> torch.Tensor:
         forward_args: Dict[str, Any] = {}
-        hash_zch_identities = self._get_hash_zch_identities(features)
-        if hash_zch_identities is not None:
+        identities_and_metadata = self._get_hash_zch_identities_and_metadata(features)
+        if identities_and_metadata is not None:
+            hash_zch_identities, hash_zch_runtime_meta = identities_and_metadata
             forward_args["hash_zch_identities"] = hash_zch_identities
+            if hash_zch_runtime_meta is not None:
+                forward_args["hash_zch_runtime_meta"] = hash_zch_runtime_meta
 
         weights = features.weights_or_none()
         if weights is not None and not torch.is_floating_point(weights):

torchrec/distributed/embedding_kernel.py

@@ -110,9 +110,9 @@ def init_raw_id_tracker(
                 get_indexed_lookups, delete
             )
 
-    def _get_hash_zch_identities(
+    def _get_hash_zch_identities_and_metadata(
         self, features: KeyedJaggedTensor
-    ) -> Optional[torch.Tensor]:
+    ) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]]:
         if self._raw_id_tracker_wrapper is None or not isinstance(
             self.emb_module, SplitTableBatchedEmbeddingBagsCodegen
         ):
@@ -131,8 +131,8 @@ def _get_hash_zch_identities(
         # across multiple training iterations. Current logic appends raw_ids from
         # all batches sequentially. This may cause misalignment with
         # features.values() which only contains the current batch.
-        raw_ids_dict = raw_id_tracker_wrapper.get_indexed_lookups(
-            table_names, emb_module.uuid
+        indexed_lookups_dict = raw_id_tracker_wrapper.get_indexed_lookups(
+            table_names, self.emb_module.uuid
         )
 
         # Build hash_zch_identities by concatenating raw IDs from tracked tables.
@@ -148,11 +148,14 @@ def _get_hash_zch_identities(
         # raw_ids are included. If some tables lack identity while others have them,
         # padding with -1 may be needed to maintain alignment.
         all_raw_ids = []
+        all_runtime_meta = []
         for table_name in table_names:
-            if table_name in raw_ids_dict:
-                raw_ids_list = raw_ids_dict[table_name]
+            if table_name in indexed_lookups_dict:
+                raw_ids_list, runtime_meta_list = indexed_lookups_dict[table_name]
                 for raw_ids in raw_ids_list:
                     all_raw_ids.append(raw_ids)
+                for runtime_meta in runtime_meta_list:
+                    all_runtime_meta.append(runtime_meta)
 
         if not all_raw_ids:
             return None
@@ -162,7 +165,12 @@ def _get_hash_zch_identities(
             f"hash_zch_identities row count ({hash_zch_identities.size(0)}) must match "
             f"features.values() length ({features.values().numel()}) to maintain 1-to-1 alignment"
         )
-        return hash_zch_identities
+
+        if all_runtime_meta:
+            hash_zch_runtime_meta = torch.cat(all_runtime_meta)
+            return (hash_zch_identities, hash_zch_runtime_meta)
+        else:
+            return (hash_zch_identities, None)
 
 
 def create_virtual_table_local_metadata(

torchrec/distributed/mc_modules.py

@@ -245,6 +245,7 @@ def __init__(
                     torch.Tensor,
                     Optional[nn.Module],
                     Optional[torch.Tensor],
+                    Optional[torch.Tensor],
                 ],
                 None,
             ]
@@ -763,13 +764,22 @@ def compute(
                         "_hash_zch_identities",
                     ):
                         if self.post_lookup_tracker_fn is not None:
+                            runtime_meta = None
+                            if (
+                                hasattr(mcm, "_hash_zch_runtime_meta")
+                                and mcm._hash_zch_runtime_meta is not None
+                            ):
+                                runtime_meta = mcm._hash_zch_runtime_meta.index_select(
+                                    dim=0, index=mc_input[table].values()
+                                )
                             self.post_lookup_tracker_fn(
                                 KeyedJaggedTensor.from_jt_dict(mc_input),
                                 torch.empty(0),
                                 None,
                                 mcm._hash_zch_identities.index_select(
                                     dim=0, index=mc_input[table].values()
                                 ),
+                                runtime_meta,
                             )
                 values = torch.cat([jt.values() for jt in output.values()])
             else:
@@ -791,11 +801,20 @@ def compute(
                 values = mc_input[table].values()
                 if hasattr(mcm, "_hash_zch_identities"):
                     if self.post_lookup_tracker_fn is not None:
+                        runtime_meta = None
+                        if (
+                            hasattr(mcm, "_hash_zch_runtime_meta")
+                            and mcm._hash_zch_runtime_meta is not None
+                        ):
+                            runtime_meta = mcm._hash_zch_runtime_meta.index_select(
+                                dim=0, index=values
+                            )
                         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),
+                            runtime_meta,
                         )
 
             remapped_kjts.append(
@@ -895,6 +914,7 @@ def register_post_lookup_tracker_fn(
                 torch.Tensor,
                 Optional[nn.Module],
                 Optional[torch.Tensor],
+                Optional[torch.Tensor],
             ],
             None,
         ],