Add support for pipelined postproc (#2940)

Summary:
Pull Request resolved: #2940

Add support for pipelined postprocs to the SparseDataDistUtil - this allows pipelined postprocs in the StagedTrainPipeline.

Differential Revision: D73824601

Author: che-sh

torchrec/distributed/train_pipeline/tests/test_train_pipelines.py

@@ -2053,7 +2053,8 @@ def gpu_postproc(x: StageOut) -> StageOut:
                 name="start_sparse_data_dist",
                 runnable=sdd.start_sparse_data_dist,
                 stream=sdd.data_dist_stream,
-                fill_callback=sdd.wait_sparse_data_dist,
+                fill_callback=sdd.wait_sdd_fill_callback,
+                data_exhausted_callback=sdd.data_exhausted_callback,
             ),
         ]
         pipeline = StagedTrainPipeline(
@@ -2120,7 +2121,8 @@ def gpu_postproc(x: StageOut) -> StageOut:
                 name="start_sparse_data_dist",
                 runnable=sdd.start_sparse_data_dist,
                 stream=sdd.data_dist_stream,
-                fill_callback=sdd.wait_sparse_data_dist,
+                fill_callback=sdd.wait_sdd_fill_callback,
+                data_exhausted_callback=sdd.data_exhausted_callback,
             ),
         ]
 
@@ -2231,7 +2233,8 @@ def test_model_detach(self) -> None:
                 name="start_sparse_data_dist",
                 runnable=sdd.start_sparse_data_dist,
                 stream=sdd.data_dist_stream,
-                fill_callback=sdd.wait_sparse_data_dist,
+                fill_callback=sdd.wait_sdd_fill_callback,
+                data_exhausted_callback=sdd.data_exhausted_callback,
             ),
         ]
 
@@ -2424,7 +2427,8 @@ def gpu_postproc(x: StageOut) -> StageOut:
                 name="start_sparse_data_dist",
                 runnable=sdd.start_sparse_data_dist,
                 stream=sdd.data_dist_stream,
-                fill_callback=sdd.wait_sparse_data_dist,
+                fill_callback=sdd.wait_sdd_fill_callback,
+                data_exhausted_callback=sdd.data_exhausted_callback,
             ),
             PipelineStage(
                 name="prefetch",

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -1624,11 +1624,9 @@ def _run_with_event(
         self,
         runnable: RunnableType,
         event: Optional[torch.Event],
-        inputs: Optional[In],
+        inputs: In,
         stream: torch.Stream,
     ) -> StageOutputWithEvent:
-        if inputs is None:
-            return (None, None)
         with self._stream_context(stream):
             # If there is no previous event, data is entering the pipeline
             if event is not None:
@@ -1672,6 +1670,11 @@ def _run_stage(
         """
         stage = self._pipeline_stages[stage_idx]
 
+        if self._debug_mode:
+            logger.info(
+                f"Running ## Pipeline Stage {stage_idx} : {stage.name} for batch {batch_offset + self._num_steps} ##",
+            )
+
         with record_function(
             f"## Pipeline Stage {stage_idx} : {stage.name} for batch {batch_offset + self._num_steps} ##"
         ):
@@ -1683,23 +1686,40 @@ def _run_stage(
                 assert batch_to_wait_with_event is not None
                 batch_to_wait, event = batch_to_wait_with_event
 
-            new_result = self._run_with_event(
-                runnable=stage.runnable,
-                event=event,
-                inputs=batch_to_wait,
-                stream=stage.stream,
-            )
+            if batch_to_wait is not None:
+                if self._debug_mode:
+                    logger.info(
+                        f"Executing ## Pipeline Stage {stage_idx} : {stage.name} for batch {batch_offset + self._num_steps} ##",
+                    )
+                new_result = self._run_with_event(
+                    runnable=stage.runnable,
+                    event=event,
+                    inputs=batch_to_wait,
+                    stream=stage.stream,
+                )
+            else:
+                if self._debug_mode:
+                    logger.info(
+                        f"Skipping due to None ## Pipeline Stage {stage_idx} : {stage.name} for batch {batch_offset + self._num_steps} ##",
+                    )
+                new_result = (None, None)
+                if (
+                    data_exhausted_callback := stage.data_exhausted_callback
+                ) is not None:
+                    data_exhausted_callback()
 
         self._stage_outputs[batch_offset] = new_result
         if self._debug_mode:
             logger.info(
-                f"Running ## Pipeline Stage {stage_idx} : {stage.name} for batch {batch_offset + self._num_steps} ##",
+                f"Finshed ## Pipeline Stage {stage_idx} : {stage.name} for batch {batch_offset + self._num_steps} ##",
             )
 
         if fill and (fill_callback := stage.fill_callback) is not None:
             if self._debug_mode:
-                logger.info(f"Finished callback for {stage.name}")
+                logger.info(f"Started callback for {stage.name}")
             fill_callback()
+            if self._debug_mode:
+                logger.info(f"Finished callback for {stage.name}")
 
         return new_result
 
@@ -1785,6 +1805,9 @@ def progress(
 
         self._num_steps += 1
 
+        if self._debug_mode:
+            logger.info(f"Starting pipeline step {self._num_steps}")
+
         for stage_idx in range(self.num_stages):
             stage_output_idx = self.num_stages - 1 - stage_idx
             self._run_stage(
@@ -1805,6 +1828,8 @@ def progress(
             self.flush_end()
             return self.progress(dataloader_iter)
 
+        if self._debug_mode:
+            logger.info(f"Finished pipeline step {self._num_steps}")
         return out