add prefetch (pytorch#3349)

Summary: Pull Request resolved: pytorch#3349

Differential Revision: D79404930

Author: peterfu0

torchrec/distributed/train_pipeline/runtime_forwards.py

@@ -58,6 +58,10 @@ def name(self) -> str:
     def args(self) -> CallArgs:
         return self._args
 
+    @classmethod
+    def prefetch(cls) -> bool:
+        return False
+
     def set_context(self, context: TForwardContext) -> None:
         self._context = context
 
@@ -220,6 +224,72 @@ def detach_embeddings(
         pass
 
 
+class PrefetchPipelinedForwardCustomizedOrder(
+    BaseForward[PrefetchTrainPipelineContext]
+):
+    """
+    This pipeline is used in TrainPipelineCustomizedOrderSparseDist
+    compute_and_output_dist for batch N is called at the end of step N - 1
+    """
+
+    def __init__(
+        self,
+        name: str,
+        args: CallArgs,
+        module: ShardedModule,
+        context: PrefetchTrainPipelineContext,
+        prefetch_stream: Optional[torch.Stream] = None,
+    ) -> None:
+        super().__init__(
+            name=name,
+            args=args,
+            module=module,
+            context=context,
+            stream=prefetch_stream,
+        )
+        self._compute_and_output_dist_awaitable: Optional[
+            Awaitable[Multistreamable]
+        ] = None
+
+    @classmethod
+    def prefetch(cls) -> bool:
+        return True
+
+    def compute_and_output_dist(self) -> None:
+        assert (
+            self._name in self._context.module_input_post_prefetch
+        ), "Invalid PrefetchPipelinedForward usage, please do not directly call model.forward()"
+        data = self._context.module_input_post_prefetch.pop(self._name)
+        ctx = self._context.module_contexts_post_prefetch.pop(self._name)
+
+        # Make sure that both result of input_dist and context
+        # are properly transferred to the current stream.
+        if self._stream is not None:
+            torch.get_device_module(self._device).current_stream().wait_stream(
+                self._stream
+            )
+            cur_stream = torch.get_device_module(self._device).current_stream()
+
+            assert isinstance(
+                data, (torch.Tensor, Multistreamable)
+            ), f"{type(data)} must implement Multistreamable interface"
+            data.record_stream(cur_stream)
+
+            ctx.record_stream(cur_stream)
+
+        self._compute_and_output_dist_awaitable = self._module.compute_and_output_dist(
+            ctx, data
+        )
+
+    # pyre-ignore [2, 24]
+    def __call__(self, *input, **kwargs) -> Awaitable:
+        if not self._compute_and_output_dist_awaitable:
+            raise Exception(
+                "compute_and_output_dist must be called before __call__",
+            )
+        return self._compute_and_output_dist_awaitable
+
+
 class PrefetchPipelinedForward(BaseForward[PrefetchTrainPipelineContext]):
     """
     This pipeline is used in PrefetchTrainPipelineSparseDist
@@ -241,6 +311,10 @@ def __init__(
             stream=prefetch_stream,
         )
 
+    @classmethod
+    def prefetch(cls) -> bool:
+        return True
+
     # pyre-ignore [2, 24]
     def __call__(self, *input, **kwargs) -> Awaitable:
         assert (

torchrec/distributed/train_pipeline/tests/test_prefetch_pipelined_forward_customized_order.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import unittest
+from unittest.mock import MagicMock
+
+from torchrec.distributed.train_pipeline.pipeline_context import (
+    PrefetchTrainPipelineContext,
+)
+from torchrec.distributed.train_pipeline.runtime_forwards import (
+    PrefetchPipelinedForwardCustomizedOrder,
+)
+from torchrec.distributed.train_pipeline.types import CallArgs
+from torchrec.distributed.train_pipeline.utils import _prefetch_enabled
+
+
+class TestPrefetchPipelinedForwardCustomizedOrder(unittest.TestCase):
+    """Test PrefetchPipelinedForwardCustomizedOrder key functionality"""
+
+    def setUp(self) -> None:
+        """Set up test fixtures."""
+        self.mock_module = MagicMock()
+        self.prefetch_context = PrefetchTrainPipelineContext()
+        self.mock_args = CallArgs(args=[], kwargs={})
+
+    def test_prefetch_returns_true(self) -> None:
+        """Test that prefetch() returns True."""
+        forward = PrefetchPipelinedForwardCustomizedOrder(
+            name="test_prefetch",
+            args=self.mock_args,
+            module=self.mock_module,
+            context=self.prefetch_context,
+        )
+
+        # Test that prefetch returns True
+        self.assertTrue(forward.prefetch())
+
+    def test_prefetch_enabled_returns_true(self) -> None:
+        """Test that _prefetch_enabled returns True for PrefetchPipelinedForwardCustomizedOrder."""
+        forward = PrefetchPipelinedForwardCustomizedOrder(
+            name="test_prefetch_enabled",
+            args=self.mock_args,
+            module=self.mock_module,
+            context=self.prefetch_context,
+        )
+
+        # Test that _prefetch_enabled returns True
+        result = _prefetch_enabled(forward)
+        self.assertTrue(result)
+
+    def test_call_fails_without_compute_and_output_dist(self) -> None:
+        """Test that __call__ fails if compute_and_output_dist is not called first."""
+        forward = PrefetchPipelinedForwardCustomizedOrder(
+            name="test_call_error",
+            args=self.mock_args,
+            module=self.mock_module,
+            context=self.prefetch_context,
+        )
+
+        # Should raise exception when called without compute_and_output_dist
+        with self.assertRaises(Exception) as context:
+            forward()
+
+        self.assertIn(
+            "compute_and_output_dist must be called before __call__",
+            str(context.exception),
+        )
+
+    def test_call_succeeds_after_compute_and_output_dist(self) -> None:
+        """Test that __call__ succeeds when compute_and_output_dist is called first."""
+        forward = PrefetchPipelinedForwardCustomizedOrder(
+            name="test_call_success",
+            args=self.mock_args,
+            module=self.mock_module,
+            context=self.prefetch_context,
+        )
+
+        # Set up mock data in context
+        test_data = MagicMock()
+        test_ctx = MagicMock()
+        self.prefetch_context.module_input_post_prefetch = {
+            "test_call_success": test_data
+        }
+        self.prefetch_context.module_contexts_post_prefetch = {
+            "test_call_success": test_ctx
+        }
+
+        # Mock the module's compute_and_output_dist method
+        mock_awaitable = MagicMock()
+        self.mock_module.compute_and_output_dist.return_value = mock_awaitable
+
+        # Call compute_and_output_dist first
+        forward.compute_and_output_dist()
+
+        # Now __call__ should succeed and return the awaitable
+        result = forward()
+        self.assertEqual(result, mock_awaitable)
+
+
+if __name__ == "__main__":
+    unittest.main()

torchrec/distributed/train_pipeline/utils.py

@@ -29,7 +29,6 @@
 
 import torch
 from torch.profiler import record_function
-
 from torchrec.distributed.dist_data import KJTAllToAll, KJTAllToAllTensorsAwaitable
 from torchrec.distributed.embedding_sharding import (
     FusedKJTListSplitsAwaitable,
@@ -53,6 +52,7 @@
     KJTAllToAllForward,
     PipelinedForward,
     PrefetchPipelinedForward,
+    PrefetchPipelinedForwardCustomizedOrder,
     TForwardContext,
 )
 from torchrec.distributed.train_pipeline.tracing import (
@@ -61,7 +61,7 @@
     Tracer,
 )
 from torchrec.distributed.train_pipeline.types import CallArgs  # noqa
-from torchrec.distributed.types import Awaitable
+from torchrec.distributed.types import Awaitable, LazyAwaitable
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 from torchrec.streamable import Multistreamable, Pipelineable
 
@@ -138,6 +138,7 @@ def _start_data_dist(
                 PrefetchPipelinedForward,
                 EmbeddingPipelinedForward,
                 InSyncEmbeddingPipelinedForward,
+                PrefetchPipelinedForwardCustomizedOrder,
             ),
         )
 
@@ -539,6 +540,10 @@ def get_next_batch(self, none_throws: bool = False) -> Optional[In]:
         return batch
 
 
+def _prefetch_enabled(forward: LazyAwaitable[Out]) -> bool:
+    return isinstance(forward, BaseForward) and forward.prefetch()
+
+
 def _prefetch_embeddings(
     batch: In,
     context: PrefetchTrainPipelineContext,
@@ -551,7 +556,11 @@ def _prefetch_embeddings(
     data_per_sharded_module = {}
     for sharded_module in pipelined_modules:
         forward = sharded_module.forward
-        assert isinstance(forward, PrefetchPipelinedForward)
+        # for backward compatibility, consider it valid if it is PrefetchPipelinedForward
+        # because the class might not have prefetch method
+        assert isinstance(forward, PrefetchPipelinedForward) or _prefetch_enabled(
+            forward
+        )
         assert forward._name in context.input_dist_tensors_requests
         request = context.input_dist_tensors_requests.pop(forward._name)
         assert isinstance(request, Awaitable)