remove duplicate modules from train pipeline and import from embedding_sharding (#1754)

Summary:
Pull Request resolved: #1754

due to issues with packaging for train pipeline in the past we needed to copy over modules into train_pipeline directly. Since the changes have since propagated through packaging we can safely remove the duplicate modules/functions.

Reviewed By: zainhuda

Differential Revision: D54497596

fbshipit-source-id: a6f532de653469dc4ed52a89ccda8e76533e4929

Author: joshuadeng

torchrec/distributed/train_pipeline/train_pipeline.py

@@ -7,11 +7,6 @@
 
 # pyre-strict
 
-"""
-NOTE: Due to an internal packaging issue, `train_pipeline.py` must be compatible with
-older versions of TorchRec. Importing new modules from other files may break model
-publishing flows.
-"""
 import abc
 import logging
 from typing import cast, Generic, Iterator, List, Optional, Tuple

torchrec/distributed/train_pipeline/utils.py

@@ -7,7 +7,6 @@
 
 # pyre-strict
 
-#!/usr/bin/env python3
 import copy
 import itertools
 import logging
@@ -34,10 +33,11 @@
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.fx.node import Node
 from torch.profiler import record_function
-from torchrec.distributed.dist_data import KJTAllToAll, KJTAllToAllTensorsAwaitable
+from torchrec.distributed.dist_data import KJTAllToAll
 from torchrec.distributed.embedding_sharding import (
-    KJTListAwaitable,
+    FusedKJTListSplitsAwaitable,
     KJTListSplitsAwaitable,
+    KJTSplitsAllToAllMeta,
 )
 from torchrec.distributed.model_parallel import DistributedModelParallel, ShardedModule
 
@@ -59,193 +59,6 @@
 StageOutputWithEvent = Tuple[Optional[StageOut], Optional[torch.cuda.Event]]
 
 
-class Tracer(torch.fx.Tracer):
-    """
-    Disables proxying buffers during tracing. Ideally, proxying buffers would be
-    disabled, but some models are currently mutating buffer values, which causes errors
-    during tracing. If those models can be rewritten to not do that, we can likely
-    remove this line.
-    """
-
-    proxy_buffer_attributes = False
-
-    def __init__(self, leaf_modules: Optional[List[str]] = None) -> None:
-        super().__init__()
-        self._leaf_modules: List[str] = leaf_modules if leaf_modules is not None else []
-
-    def is_leaf_module(self, m: torch.nn.Module, module_qualified_name: str) -> bool:
-        if (
-            isinstance(m, ShardedModule)
-            or module_qualified_name in self._leaf_modules
-            or isinstance(m, FSDP)
-        ):
-            return True
-        return super().is_leaf_module(m, module_qualified_name)
-
-
-# TODO: remove after packaging issue is resolved.
-class SplitsAllToAllAwaitable(Awaitable[List[List[int]]]):
-    def __init__(
-        self,
-        input_tensors: List[torch.Tensor],
-        pg: dist.ProcessGroup,
-    ) -> None:
-        super().__init__()
-        self.num_workers: int = pg.size()
-
-        with record_function("## all2all_data:kjt splits ##"):
-            self._output_tensor: torch.Tensor = torch.empty(
-                [self.num_workers * len(input_tensors)],
-                device=input_tensors[0].device,
-                dtype=input_tensors[0].dtype,
-            )
-            input_tensor = torch.stack(input_tensors, dim=1).flatten()
-            self._splits_awaitable: dist.Work = dist.all_to_all_single(
-                output=self._output_tensor,
-                input=input_tensor,
-                group=pg,
-                async_op=True,
-            )
-
-    def _wait_impl(self) -> List[List[int]]:
-        self._splits_awaitable.wait()
-        return self._output_tensor.view(self.num_workers, -1).T.tolist()
-
-
-# TODO: remove after packaging issue is resolved.
-C = TypeVar("C", bound=Multistreamable)
-T = TypeVar("T")
-
-
-# TODO: remove after packaging issue is resolved.
-def _set_sharding_context_intra_a2a(
-    tensors_awaitables: List[Awaitable[KeyedJaggedTensor]],
-    ctx: C,
-) -> None:
-    for awaitable, sharding_context in zip(
-        tensors_awaitables,
-        getattr(ctx, "sharding_contexts", []),
-    ):
-        if isinstance(awaitable, KJTAllToAllTensorsAwaitable):
-            if hasattr(sharding_context, "input_splits"):
-                sharding_context.input_splits = awaitable._input_splits["values"]
-            if hasattr(sharding_context, "output_splits"):
-                sharding_context.output_splits = awaitable._output_splits["values"]
-            if hasattr(sharding_context, "sparse_features_recat"):
-                sharding_context.sparse_features_recat = awaitable._recat
-            if (
-                hasattr(sharding_context, "batch_size_per_rank")
-                and awaitable._stride_per_rank is not None
-            ):
-                sharding_context.batch_size_per_rank = awaitable._stride_per_rank
-
-
-# TODO: remove after packaging issue is resolved.
-@dataclass
-class KJTSplitsAllToAllMeta:
-    pg: dist.ProcessGroup
-    _input: KeyedJaggedTensor
-    splits: List[int]
-    splits_tensors: List[torch.Tensor]
-    input_splits: List[List[int]]
-    input_tensors: List[torch.Tensor]
-    labels: List[str]
-    keys: List[str]
-    device: torch.device
-    stagger: int
-
-
-# TODO: remove after packaging issue is resolved.
-def _split(flat_list: List[T], splits: List[int]) -> List[List[T]]:
-    return [
-        flat_list[sum(splits[:i]) : sum(splits[:i]) + n] for i, n in enumerate(splits)
-    ]
-
-
-# TODO: remove after packaging issue is resolved.
-class FusedKJTListSplitsAwaitable(Awaitable[List[KJTListAwaitable]]):
-    def __init__(
-        self,
-        requests: List[KJTListSplitsAwaitable[C]],
-        contexts: List[C],
-        pg: Optional[dist.ProcessGroup],
-    ) -> None:
-        super().__init__()
-        self._contexts = contexts
-        self._awaitables: List[
-            Union[KJTSplitsAllToAllMeta, Awaitable[Awaitable[KeyedJaggedTensor]]]
-        ] = [awaitable for request in requests for awaitable in request.awaitables]
-        self._output_lengths: List[int] = [
-            len(request.awaitables) for request in requests
-        ]
-        self._lengths: List[int] = [
-            (
-                len(awaitable.splits_tensors)
-                if isinstance(awaitable, KJTSplitsAllToAllMeta)
-                else 0
-            )
-            for awaitable in self._awaitables
-        ]
-        splits_tensors = [
-            splits_tensor
-            for awaitable in self._awaitables
-            for splits_tensor in (
-                awaitable.splits_tensors
-                if isinstance(awaitable, KJTSplitsAllToAllMeta)
-                else []
-            )
-        ]
-        self._splits_awaitable: Optional[SplitsAllToAllAwaitable] = (
-            SplitsAllToAllAwaitable(
-                input_tensors=splits_tensors,
-                pg=pg,
-            )
-            if splits_tensors and pg is not None
-            else None
-        )
-
-    def _wait_impl(self) -> List[KJTListAwaitable]:
-        if self._splits_awaitable:
-            splits_list = self._splits_awaitable.wait()
-            splits_per_awaitable = _split(splits_list, self._lengths)
-        else:
-            splits_per_awaitable = [[] for _ in range(len(self._lengths))]
-        tensors_awaitables = []
-        for splits, awaitable in zip(splits_per_awaitable, self._awaitables):
-            if not splits:  # NoWait
-                assert isinstance(awaitable, Awaitable)
-                tensors_awaitables.append(awaitable.wait())
-                continue
-            assert isinstance(awaitable, KJTSplitsAllToAllMeta)
-            if awaitable._input.variable_stride_per_key():
-                output_splits = splits
-                stride_per_rank = None
-            else:
-                output_splits = splits[:-1]
-                stride_per_rank = splits[-1]
-            tensors_awaitables.append(
-                KJTAllToAllTensorsAwaitable(
-                    pg=awaitable.pg,
-                    input=awaitable._input,
-                    splits=awaitable.splits,
-                    input_splits=awaitable.input_splits,
-                    output_splits=output_splits,
-                    input_tensors=awaitable.input_tensors,
-                    labels=awaitable.labels,
-                    keys=awaitable.keys,
-                    device=awaitable.device,
-                    stagger=awaitable.stagger,
-                    stride_per_rank=stride_per_rank,
-                )
-            )
-        output = []
-        awaitables_per_output = _split(tensors_awaitables, self._output_lengths)
-        for awaitables, ctx in zip(awaitables_per_output, self._contexts):
-            _set_sharding_context_intra_a2a(awaitables, ctx)
-            output.append(KJTListAwaitable(awaitables, ctx))
-        return output
-
-
 @dataclass
 class TrainPipelineContext:
     """
@@ -462,6 +275,30 @@ def __call__(self, input: KeyedJaggedTensor) -> KJTSplitsAllToAllMeta:
             )
 
 
+class Tracer(torch.fx.Tracer):
+    """
+    Disables proxying buffers during tracing. Ideally, proxying buffers would be
+    disabled, but some models are currently mutating buffer values, which causes errors
+    during tracing. If those models can be rewritten to not do that, we can likely
+    remove this line.
+    """
+
+    proxy_buffer_attributes = False
+
+    def __init__(self, leaf_modules: Optional[List[str]] = None) -> None:
+        super().__init__()
+        self._leaf_modules: List[str] = leaf_modules if leaf_modules is not None else []
+
+    def is_leaf_module(self, m: torch.nn.Module, module_qualified_name: str) -> bool:
+        if (
+            isinstance(m, ShardedModule)
+            or module_qualified_name in self._leaf_modules
+            or isinstance(m, FSDP)
+        ):
+            return True
+        return super().is_leaf_module(m, module_qualified_name)
+
+
 def _to_device(batch: In, device: torch.device, non_blocking: bool) -> In:
     assert isinstance(
         batch, (torch.Tensor, Pipelineable)