Add a option to move embedding lookup after sparse data dist in FusedSDD (#3345)

Summary:
Pull Request resolved: #3345

This diff adds an option to allows the embedding lookup trigger after the sparse data dist. This can potentially improve performance when CPU is blocked by sparse data dist kernel launch and could not launch forward kernel earlier.
 {F1981658737,width=300}

Reviewed By: TroyGarden

Differential Revision: D81494775

fbshipit-source-id: 2d975d6784b806edbcdfdf44a181632570e2c940

Author: Jingchang Zhang

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -964,6 +964,7 @@ def __init__(
         ] = None,
         strict: bool = False,
         emb_lookup_stream: str = "data_dist",  # new, current, data_dist (default)
+        embedding_lookup_after_data_dist: bool = False,
     ) -> None:
         super().__init__(
             model=model,
@@ -975,6 +976,8 @@ def __init__(
             pipeline_postproc=pipeline_postproc,
             custom_model_fwd=custom_model_fwd,
         )
+        self._embedding_lookup_after_data_dist = embedding_lookup_after_data_dist
+
         if emb_lookup_stream == "new":
             self._emb_lookup_stream: Optional[torch.Stream] = (
                 (torch.get_device_module(device).Stream())
@@ -1046,8 +1049,9 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
         self._set_module_context(self.contexts[0])
 
         # start embedding_lookup so it can overlap with previous optimizer
-        # pyre-ignore [6]
-        self.start_embedding_lookup(self.batches[0], self.contexts[0])
+        if not self._embedding_lookup_after_data_dist:
+            # pyre-ignore [6]
+            self.start_embedding_lookup(self.batches[0], self.contexts[0])
 
         if self._model.training:
             with record_function("## zero_grad ##"):
@@ -1064,6 +1068,10 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
         # batch i+2: load data and copy to gpu, the dataload iter will first exhaust here
         self.enqueue_batch(dataloader_iter)
 
+        if self._embedding_lookup_after_data_dist:
+            # pyre-ignore [6]
+            self.start_embedding_lookup(self.batches[0], self.contexts[0])
+
         # forward
         with record_function(f"## forward {self.contexts[0].index} ##"):
             losses, output = self._model_fwd(self.batches[0])