Revert "[Traceable FSDP2][Inductor] Re-inplace all_gather_into_tensor (#129773)"

pytorchmergebot · pytorchmergebot · commit c1e7e40f246e · 2024-07-16T20:54:14.000Z
This reverts commit f2f3102. Reverted #129773 on behalf of https://github.com/clee2000 due to failed inductor/test_torchinductor_dynamic_shapes.py on mac https://github.com/pytorch/pytorch/actions/runs/9963396991/job/27530249256 https://hud.pytorch.org/pytorch/pytorch/commit/f2f31027ce8dc3985663bf6eaa66f3c5559b724a. The build failed on PR so test jobs didn't run ([comment](#129773 (comment)))
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_compile.py b/test/distributed/_composable/fsdp/test_fully_shard_compile.py
@@ -3,23 +3,18 @@
 
 import contextlib
 import copy
-import functools
 import unittest
-from unittest import mock
 
 import torch
 import torch._dynamo.testing
 import torch.distributed._composable.fsdp._fsdp_param
 from torch import nn
 from torch._dynamo import compiled_autograd
-from torch._inductor import comms
-from torch._inductor.utils import run_and_get_code
 
 from torch.distributed._composable.fsdp import fully_shard
 from torch.distributed._composable.fsdp._fsdp_common import TrainingState
 from torch.distributed._composable.fsdp._fsdp_param_group import FSDPParamGroup
 from torch.distributed._tensor import init_device_mesh
-from torch.testing import FileCheck
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import FSDPTest, MLP
 from torch.testing._internal.common_utils import run_tests, skipIfRocm
@@ -30,10 +25,6 @@
 from torch.utils._triton import has_triton
 
 
-def _is_op_in_graph(graph, op):
-    return any(node.target is op for node in graph.nodes)
-
-
 class TestFullyShardCompileCompute(FSDPTest):
     @unittest.skipIf(not has_triton(), "Inductor+gpu needs triton and recent GPU arch")
     @skip_if_lt_x_gpu(2)
@@ -142,27 +133,6 @@ def f(x):
         torch.compile(f, backend="aot_eager")(x)
         self.assertEqual(x, ref_x)
 
-    def _reinplace_all_gather_with_checks(self, graph, orig_fn):
-        self.assertTrue(
-            _is_op_in_graph(
-                graph,
-                torch.ops._c10d_functional.all_gather_into_tensor.default,
-            )
-        )
-        orig_fn(graph)
-        self.assertFalse(
-            _is_op_in_graph(
-                graph,
-                torch.ops._c10d_functional.all_gather_into_tensor.default,
-            )
-        )
-        self.assertTrue(
-            _is_op_in_graph(
-                graph,
-                torch.ops._c10d_functional.all_gather_into_tensor_out.default,
-            )
-        )
-
     @torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
     @torch._functorch.config.patch(recompute_views=True)
     @torch._functorch.config.patch(cse=False)
@@ -328,27 +298,9 @@ def test_transformer_fullgraph_backend_aot_eager_decomp_partition(self):
     # TODO: native_dropout causes CUDA IMA error, need to figure out why
     @torch._inductor.config.patch(fallback_random=True)
     def test_transformer_fullgraph_backend_inductor(self):
-        with mock.patch.object(
-            comms,
-            "reinplace_fsdp_all_gather",
-            functools.partial(
-                self._reinplace_all_gather_with_checks,
-                orig_fn=comms.reinplace_fsdp_all_gather,
-            ),
-        ):
-            _, triton_codes = run_and_get_code(
-                lambda: self._test_traceable_fsdp(
-                    *self._create_transformer_factory_fns(), "inductor", fullgraph=True
-                )
-            )
-        self.assertTrue(
-            len(triton_codes) == 2,
-            "Expected two separate lowerings to Triton code, one from FWD graph and one from Compiled Autograd BWD graph",
+        self._test_traceable_fsdp(
+            *self._create_transformer_factory_fns(), "inductor", fullgraph=True
         )
-        for code in triton_codes:
-            FileCheck().check(
-                "torch.ops._c10d_functional.all_gather_into_tensor_out."
-            ).run(code)
 
 
 if __name__ == "__main__":
diff --git a/torch/_inductor/comms.py b/torch/_inductor/comms.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import heapq
-import operator
 
 import sys
 from collections import defaultdict
@@ -319,102 +318,3 @@ def reorder_compute_and_comm_for_overlap(
             except Exception as e:
                 overlap_log.debug(str(e))
     return order
-
-
-def reinplace_fsdp_all_gather(graph: torch.fx.Graph) -> None:
-    from .pattern_matcher import (
-        CallFunction,
-        KeywordArg,
-        Match,
-        PatternMatcherPass,
-        register_graph_pattern,
-    )
-
-    """
-    all_gather_copy_in = torch.ops.fsdp.all_gather_copy_in.default(...);
-    getitem = all_gather_copy_in[0];
-    (getitem_1 = all_gather_copy_in[1];)  # optional
-
-    all_gather_into_tensor = torch.ops._c10d_functional.all_gather_into_tensor.default(getitem, ...);
-
-    ->
-
-    all_gather_copy_in = torch.ops.fsdp.all_gather_copy_in.default(...);
-    getitem = all_gather_copy_in[0];
-    getitem_1 = all_gather_copy_in[1];
-
-    all_gather_into_tensor = torch.ops._c10d_functional.all_gather_into_tensor_out.default(getitem, ..., out=getitem_1);
-    """
-
-    def remove_unused_getitem(g):
-        # Remove `getitem_X = all_gather_copy_in[1]` which is never used.
-        node_list = list(g.nodes)
-        for n in node_list:
-            if (
-                n.target == operator.getitem
-                and n.args[0].target is torch.ops.fsdp.all_gather_copy_in.default
-                and n.args[1] == 1
-            ):
-                g.erase_node(n)
-
-    graph_pass = PatternMatcherPass()
-
-    @register_graph_pattern(
-        CallFunction(
-            torch.ops._c10d_functional.all_gather_into_tensor.default,
-            CallFunction(
-                operator.getitem,
-                CallFunction(
-                    torch.ops.fsdp.all_gather_copy_in.default,
-                    KeywordArg("all_gather_inputs"),
-                    KeywordArg("inp_split_sizes"),
-                    KeywordArg("all_gather_input_numel"),
-                    KeywordArg("world_size"),
-                    KeywordArg("rank"),
-                    KeywordArg("dtype"),
-                    KeywordArg("device"),
-                ),
-                KeywordArg("item_idx"),
-            ),
-            KeywordArg("group_size"),
-            KeywordArg("group_name"),
-        ),
-        pass_dict=graph_pass,
-        extra_check=lambda match: match.kwargs["item_idx"] == 0,
-    )
-    def reinplace_all_gather(match: Match, *args, **kwargs):
-        def repl(
-            *args,
-        ):
-            copy_in_args = args[:-2]
-            group_size = args[-2]
-            group_name = args[-1]
-            all_gather_copy_in = torch.ops.fsdp.all_gather_copy_in.default(
-                *copy_in_args
-            )
-            getitem = all_gather_copy_in[0]
-            getitem_1 = all_gather_copy_in[1]
-            all_gather_into_tensor = (
-                torch.ops._c10d_functional.all_gather_into_tensor_out.default(
-                    getitem, group_size, group_name, out=getitem_1
-                )
-            )
-            return all_gather_into_tensor
-
-        match.replace_by_example(
-            repl,
-            [
-                kwargs["all_gather_inputs"],
-                kwargs["inp_split_sizes"],
-                kwargs["all_gather_input_numel"],
-                kwargs["world_size"],
-                kwargs["rank"],
-                kwargs["dtype"],
-                kwargs["device"],
-                kwargs["group_size"],
-                kwargs["group_name"],
-            ],
-        )
-
-    remove_unused_getitem(graph)
-    graph_pass.apply(graph)  # type: ignore[arg-type]
diff --git a/torch/_inductor/fx_passes/post_grad.py b/torch/_inductor/fx_passes/post_grad.py
@@ -12,7 +12,6 @@
 from torch import fx
 from torch._decomp import register_decomposition
 from torch._dynamo.utils import counters, optimus_scuba_log
-from torch._inductor import comms
 from torch._inductor.virtualized import ops
 
 from torch._prims_common import is_boolean_dtype, is_expandable_to, is_integer_dtype
@@ -142,8 +141,6 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
     reinplace_inplaceable_ops(gm.graph)
     decompose_auto_functionalized(gm.graph)
 
-    comms.reinplace_fsdp_all_gather(gm.graph)
-
     gm.recompile()
     optimus_scuba_log["after_recompile_post_grad"] = upload_graph(gm.graph)
     gm.graph.lint()
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -6798,14 +6798,7 @@ def create_inplace(
         packed.cpp_kernel_name = cpp_kernel_name
         packed.python_kernel_name = python_kernel_name
 
-        inps = pytree.tree_leaves(inputs)
-        mark_node_as_mutating(packed, *inps)
-        # For inplace collective ops, the input is guaranteed to be alias of the returned value of op.
-        packed.alias_names.extend([inp.get_name() for inp in inps])
-        if "out" in kwargs:
-            mark_node_as_mutating(packed, kwargs["out"])
-            # For out-variant collective ops, the `out=` arg is guaranteed to be alias of the returned value of op.
-            packed.alias_names.append(kwargs["out"].get_name())
+        mark_node_as_mutating(packed, *pytree.tree_leaves(inputs))
 
     # NOTE: [Out-of-Place Collective Safety]
     # Between the initiation and completion of an out-of-place collective:
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -207,14 +207,6 @@ def get_overloads(aten_fn):
     return aten_fn
 
 
-def in_namespace(op, namespace):
-    if isinstance(op, torch._ops.OpOverloadPacket):
-        return namespace in op._qualified_op_name
-    elif isinstance(op, torch._ops.OpOverload):
-        return namespace in op.name()
-    return False
-
-
 def transform_args(args, broadcast, type_promotion_kind, convert_input_to_bool):
     indices = [i for i, x in enumerate(args) if isinstance(x, TensorBox)]
     if (type_promotion_kind or convert_input_to_bool) and indices:
@@ -301,9 +293,7 @@ def wrapped(*args, **kwargs):
             args = args[0]
 
         # kwargs tensors not supported yet unless it's a fallback op
-        if not all(
-            (fn in fallbacks or in_namespace(fn, "_c10d_functional")) for fn in aten_fn
-        ):
+        if not all(fn in fallbacks for fn in aten_fn):
             assert not any(isinstance(x, TensorBox) for x in kwargs.values())
             # explicitly assert for "out=" ops for better error messages
             assert not any(
@@ -6376,17 +6366,6 @@ def _all_gather_into_tensor_coalesced(inputs, group_size, group_name):
             ),
         )
 
-    @register_lowering(_c10d_functional.all_gather_into_tensor_out)
-    def _all_gather_into_tensor_out(inp, group_size, group_name, *, out):
-        ir._CollectiveKernel.create_inplace(
-            _c10d_functional.all_gather_into_tensor_out.default,
-            inp,
-            group_size,
-            group_name,
-            out=out,
-        )
-        return out
-
     @register_lowering(_c10d_functional.reduce_scatter_tensor)
     def _reduce_scatter_tensor(inp, reduce_op, group_size, group_name):
         return ir.TensorBox.create(
diff --git a/torch/_inductor/utils.py b/torch/_inductor/utils.py
@@ -1856,7 +1856,7 @@ def extract_tensor_metadata(input: torch.Tensor) -> Dict[str, Any]:
                 metadata["sizes"] = list(input.size())
                 metadata["strides"] = list(input.stride())
                 metadata["requires_grad"] = input.requires_grad
-                metadata["dispatch_key_set"] = torch._C._dispatch_keys(input).raw_repr()  # type: ignore[attr-defined]
+                metadata["dispatch_key_set"] = torch._C._dispatch_keys(input).raw_repr()
                 return metadata
 
             def extract_scalar_metadata(
