Updating record_lookup function signature to accommodate future implementations (#3471)

Summary:
Pull Request resolved: #3471

internal

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

Specific Context: Update record_lookup function signature to accommodate MRS DeltaTracker implementation

Differential Revision: D80614980

fbshipit-source-id: 84f874668f7f5a5916611c93e58577d6f5dc00bc

Author: aliafzal

torchrec/distributed/embedding.py

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

torchrec/distributed/embedding_lookup.py

@@ -210,7 +210,7 @@ def __init__(
         self.grouped_configs = grouped_configs
         # Model tracker function to tracker optimizer state
         self.optim_state_tracker_fn: Optional[
-            Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]
+            Callable[[KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None]
         ] = None
 
     def _create_embedding_kernel(
@@ -325,7 +325,7 @@ def forward(
             # Model tracker optimizer state function, will only be set called
             # when model tracker is configured to track optimizer state
             if self.optim_state_tracker_fn is not None:
-                self.optim_state_tracker_fn(emb_op, features, lookup)
+                self.optim_state_tracker_fn(features, lookup, emb_op)
 
         return embeddings_cat_empty_rank_handle(embeddings, self._dummy_embs_tensor)
 
@@ -432,13 +432,15 @@ def purge(self) -> None:
 
     def register_optim_state_tracker_fn(
         self,
-        record_fn: Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None],
+        record_fn: Callable[
+            [KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None
+        ],
     ) -> None:
         """
         Model tracker function to tracker optimizer state
 
          Args:
-             record_fn (Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]): A custom record function to be called after lookup is done.
+             record_fn (Callable[[KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None]): A custom record function to be called after lookup is done.
 
         """
         self.optim_state_tracker_fn = record_fn
@@ -544,7 +546,7 @@ def __init__(
         )
         # Model tracker function to tracker optimizer state
         self.optim_state_tracker_fn: Optional[
-            Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]
+            Callable[[KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None]
         ] = None
 
     def _create_embedding_kernel(
@@ -710,7 +712,7 @@ def forward(
                 # Model tracker optimizer state function, will only be set called
                 # when model tracker is configured to track optimizer state
                 if self.optim_state_tracker_fn is not None:
-                    self.optim_state_tracker_fn(emb_op, features, lookup)
+                    self.optim_state_tracker_fn(features, lookup, emb_op)
 
                 if features.variable_stride_per_key() and len(self._emb_modules) > 1:
                     stride_per_rank_per_key = list(
@@ -845,13 +847,15 @@ def purge(self) -> None:
 
     def register_optim_state_tracker_fn(
         self,
-        record_fn: Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None],
+        record_fn: Callable[
+            [KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None
+        ],
     ) -> None:
         """
         Model tracker function to tracker optimizer state
 
          Args:
-             record_fn (Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]): A custom record function to be called after lookup is done.
+             record_fn (Callable[[KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None]): A custom record function to be called after lookup is done.
 
         """
         self.optim_state_tracker_fn = record_fn

torchrec/distributed/embedding_types.py

@@ -391,7 +391,7 @@ def __init__(
         self._lookups: List[nn.Module] = []
         self._output_dists: List[nn.Module] = []
         self.post_lookup_tracker_fn: Optional[
-            Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]
+            Callable[[KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None]
         ] = None
         self.post_odist_tracker_fn: Optional[Callable[..., None]] = None
 
@@ -444,14 +444,16 @@ def train(self, mode: bool = True):  # pyre-ignore[3]
 
     def register_post_lookup_tracker_fn(
         self,
-        record_fn: Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None],
+        record_fn: Callable[
+            [KeyedJaggedTensor, torch.Tensor, Optional[nn.Module]], None
+        ],
     ) -> None:
         """
         Register a function 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.
+            record_fn (Callable[[KeyedJaggedTensor, torch.Tensor,Optional[nn.Module]], None]): A custom record function to be called after lookup is done.
 
         """
         if self.post_lookup_tracker_fn is not 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(self, features, embs)
+                    self.post_lookup_tracker_fn(features, embs, self)
 
             with maybe_annotate_embedding_event(
                 EmbeddingEvent.OUTPUT_DIST,

torchrec/distributed/model_tracker/model_delta_tracker.py

@@ -79,7 +79,10 @@ class ModelDeltaTracker(ABC):
 
     @abstractmethod
     def record_lookup(
-        self, emb_module: nn.Module, kjt: KeyedJaggedTensor, states: torch.Tensor
+        self,
+        kjt: KeyedJaggedTensor,
+        states: torch.Tensor,
+        emb_module: Optional[nn.Module] = None,
     ) -> None:
         """
         Records the IDs from a given KeyedJaggedTensor and their corresponding embeddings/parameter states.
@@ -233,7 +236,10 @@ def trigger_compaction(self) -> None:
             self.curr_compact_index = end_idx
 
     def record_lookup(
-        self, emb_module: nn.Module, kjt: KeyedJaggedTensor, states: torch.Tensor
+        self,
+        kjt: KeyedJaggedTensor,
+        states: torch.Tensor,
+        emb_module: Optional[nn.Module] = None,
     ) -> None:
         """
         Records the IDs from a given KeyedJaggedTensor and their corresponding embeddings/parameter states.
@@ -258,12 +264,12 @@ def record_lookup(
             self.record_embeddings(kjt, states)
         # In MOMENTUM_LAST mode, we track per feature IDs and corresponding momentum values received in the current batch.
         elif self._mode == TrackingMode.MOMENTUM_LAST:
-            self.record_momentum(emb_module, kjt)
+            self.record_momentum(none_throws(emb_module), kjt)
         elif (
             self._mode == TrackingMode.MOMENTUM_DIFF
             or self._mode == TrackingMode.ROWWISE_ADAGRAD
         ):
-            self.record_rowwise_optim_state(emb_module, kjt)
+            self.record_rowwise_optim_state(none_throws(emb_module), kjt)
         else:
             raise NotImplementedError(f"Tracking mode {self._mode} is not supported")