Update on "Make FakeProcessGroup traceable"

fegin · fegin · commit 858ef9123d95 · 2023-11-08T17:09:23.000-08:00
This PR mimics what we have done to trace ProcessGroup. This allows use to use FakeProcessGroup with torch.compile. FakeProcessGroup allows us to use world_size > 1 without creating multiple processes thus enabling the usage of PDB to debug bucketing DDP allreduce in the Inductor. We can theoretically use GLOO with world_size==1 to achieve the same goal. However, the `wait()` seems to be optimized away when the world_size is 1. Differential Revision: [D51136463](https://our.internmc.facebook.com/intern/diff/D51136463/) cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -92,7 +92,6 @@
 )
 from .distributed import (
     DeviceMeshVariable,
-    FakeProcessGroupVariable,
     PlacementClassVariable,
     PlacementVariable,
     ProcessGroupVariable,
@@ -661,12 +660,6 @@ def index_source(key):
                 source=self.source,
                 guards=self.make_guards(GuardBuilder.ID_MATCH),
             )
-        elif FakeProcessGroupVariable.is_process_group(value):
-            return FakeProcessGroupVariable(
-                value,
-                source=self.source,
-                guards=self.make_guards(GuardBuilder.ID_MATCH),
-            )
         elif DeviceMeshVariable.is_device_mesh(value):
             # TODO: see if we need to add custom guard instead
             # of a simple ID_MATCH
diff --git a/torch/_dynamo/variables/distributed.py b/torch/_dynamo/variables/distributed.py
@@ -211,16 +211,6 @@ def is_process_group(value):
         if not DistributedVariable.is_available():
             return False
         from torch._C._distributed_c10d import ProcessGroup
-
-        return istype(value, ProcessGroup)
-
-
-class FakeProcessGroupVariable(ProcessGroupVariable):
-    @staticmethod
-    def is_process_group(value):
-        if not DistributedVariable.is_available():
-            return False
-
         from torch.testing._internal.distributed.fake_pg import FakeProcessGroup
 
-        return istype(value, FakeProcessGroup)
+        return istype(value, (ProcessGroup, FakeProcessGroup))
diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
@@ -42,12 +42,7 @@
     NullContextVariable,
     TorchFunctionDisableVariable,
 )
-from .distributed import (
-    FakeProcessGroupVariable,
-    is_constant_pg_functions,
-    is_from_local,
-    ProcessGroupVariable,
-)
+from .distributed import is_constant_pg_functions, is_from_local, ProcessGroupVariable
 from .higher_order_ops import TorchHigherOrderOperatorVariable
 from .lists import ListVariable, TupleVariable
 from .torch_function import can_dispatch_torch_function, dispatch_torch_function
@@ -589,7 +584,7 @@ def call_function(
             # We desugar it at trace-time into ranks by directly calling util
             # bake the result into the trace
             assert len(args) == 1, "Expected one arg (pg)"
-            assert isinstance(args[0], (ProcessGroupVariable, FakeProcessGroupVariable))
+            assert isinstance(args[0], ProcessGroupVariable)
 
             invocation_result = self.value(args[0].as_python_constant())
             # Note - while we *could* cook up sources around invocations, like a FunctionSource
