- Add CompiledAutograd pipeline

Summary: Add new pipeline for CompiledAutograd development.

Differential Revision: D61403499

Author: flaviotruzzi

torchrec/distributed/train_pipeline/__init__.py

@@ -16,6 +16,7 @@
     TrainPipelineBase,  # noqa
     TrainPipelinePT2,  # noqa
     TrainPipelineSparseDist,  # noqa
+    TrainPipelineSparseDistCompAutograd,  # noqa
 )
 from torchrec.distributed.train_pipeline.utils import (  # noqa
     _override_input_dist_forwards,  # noqa

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -8,13 +8,15 @@
 # pyre-strict
 
 import abc
+import contextlib
 import logging
 from collections import deque
 from dataclasses import dataclass
 from typing import (
     Any,
     Callable,
     cast,
+    ContextManager,
     Deque,
     Dict,
     Generic,
@@ -28,6 +30,7 @@
 
 import torch
 from torch.autograd.profiler import record_function
+from torchrec.distributed.comm_ops import set_use_sync_collectives
 from torchrec.distributed.dist_data import KJTAllToAllTensorsAwaitable
 from torchrec.distributed.model_parallel import ShardedModule
 from torchrec.distributed.train_pipeline.utils import (
@@ -1506,3 +1509,63 @@ def progress(
             return self.progress(dataloader_iter)
 
         return out
+
+
+class TrainPipelineSparseDistCompAutograd(TrainPipelineSparseDist[In, Out]):
+    """
+    This pipeline clone the TrainPipelineSparseDist, but execute the progress
+    method within compiled autograd context.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        device: torch.device,
+        execute_all_batches: bool = True,
+        apply_jit: bool = False,
+        context_type: Type[TrainPipelineContext] = TrainPipelineContext,
+        pipeline_preproc: bool = False,
+        custom_model_fwd: Optional[
+            Callable[[In], Tuple[torch.Tensor, List[torch.Tensor]]]
+        ] = None,
+    ) -> None:
+        set_use_sync_collectives(True)
+        super().__init__(
+            model,
+            optimizer,
+            device,
+            execute_all_batches,
+            apply_jit,
+            context_type,
+            pipeline_preproc,
+            custom_model_fwd,
+        )
+
+    @staticmethod
+    def get_compiled_autograd_ctx(
+        model: torch.nn.Module,
+    ) -> ContextManager:
+        compiled_autograd = (
+            hasattr(model, "_compiled_autograd") and model._compiled_autograd
+        )
+
+        model._compiled_autograd_options = {
+            "backend": "inductor",
+            "dynamic": True,
+            "fullgraph": True,
+        }
+        torch._dynamo.config.optimize_ddp = "python_reducer"
+        return (
+            torch._dynamo.compiled_autograd.enable(
+                torch.compile(**model._compiled_autograd_options)
+            )
+            if compiled_autograd
+            else contextlib.nullcontext()
+        )
+
+    def progress(self, dataloader_iter: Iterator[In]) -> Out:
+        self._model._compiled_autograd = True
+
+        with self.get_compiled_autograd_ctx(self._model):
+            return super().progress(dataloader_iter)