Revert "Switch to native functional collective by default (#120370)"

pytorchmergebot · pytorchmergebot · commit 41adec3c59ce · 2024-02-26T21:55:13.000Z
This reverts commit 1f1bc0e. Reverted #120370 on behalf of https://github.com/yifuwang due to broke CI ([comment](#120370 (comment)))
diff --git a/test/distributed/_tensor/experimental/test_tp_transform.py b/test/distributed/_tensor/experimental/test_tp_transform.py
@@ -92,8 +92,8 @@ def test_tp_transform_with_uncovered_op(self):
         self.assert_has_c10d_ops(
             tp_exported_program.graph_module,
             {
-                "_c10d_functional.all_gather_into_tensor.default": 1,
-                "_c10d_functional.wait_tensor.default": 1,
+                "c10d_functional.all_gather_into_tensor.default": 1,
+                "c10d_functional.wait_tensor.default": 1,
             },
         )
 
@@ -131,8 +131,8 @@ def test_tp_transform_e2e(self):
         self.assert_has_c10d_ops(
             tp_exported_program.graph_module,
             {
-                "_c10d_functional.all_reduce.default": 2,
-                "_c10d_functional.wait_tensor.default": 2,
+                "c10d_functional.all_reduce.default": 2,
+                "c10d_functional.wait_tensor.default": 2,
             },
         )
 
@@ -167,8 +167,8 @@ def test_tp_transform_no_bias(self):
         self.assert_has_c10d_ops(
             tp_exported_program.graph_module,
             {
-                "_c10d_functional.all_reduce.default": 1,
-                "_c10d_functional.wait_tensor.default": 1,
+                "c10d_functional.all_reduce.default": 1,
+                "c10d_functional.wait_tensor.default": 1,
             },
         )
 
diff --git a/torch/distributed/_functional_collectives_impl.py b/torch/distributed/_functional_collectives_impl.py
@@ -7,7 +7,6 @@
 import torch
 import torch.distributed as dist
 import torch.distributed.distributed_c10d as c10d
-from torch._dynamo import assume_constant_result
 
 """
 Moved eager kernel implementations to a separate file partly for readability and partly as it is currently
@@ -25,23 +24,10 @@
 
 """
 
-_use_native_funcol: Optional[bool] = None
+_use_native_funcol = "_USE_NATIVE_C10D_FUNCTIONAL" in os.environ
 
 
-@assume_constant_result
 def native_funcol_enabled():
-    global _use_native_funcol
-    if _use_native_funcol is None:
-        try:
-            # Disable native funcol when torch_xla is installed. This check
-            # will be removed once torch_xla adopts the native_funcol IR.
-            import torch_xla  # noqa: F401
-
-            _use_native_funcol = False
-        except Exception:
-            # When TORCH_DISABLE_NATIVE_FUNCOL is set, fallback to py funcol
-            _use_native_funcol = os.environ.get("TORCH_DISABLE_NATIVE_FUNCOL") != "1"
-
     return _use_native_funcol
 
 

Original file line number	Diff line number	Diff line change
`@@ -92,8 +92,8 @@ def test_tp_transform_with_uncovered_op(self):`
`92`	`92`	`self.assert_has_c10d_ops(`
`93`	`93`	`tp_exported_program.graph_module,`
`94`	`94`	`{`
`95`		`- "_c10d_functional.all_gather_into_tensor.default": 1,`
`96`		`- "_c10d_functional.wait_tensor.default": 1,`
	`95`	`+ "c10d_functional.all_gather_into_tensor.default": 1,`
	`96`	`+ "c10d_functional.wait_tensor.default": 1,`
`97`	`97`	`},`
`98`	`98`	`)`
`99`	`99`
`@@ -131,8 +131,8 @@ def test_tp_transform_e2e(self):`
`131`	`131`	`self.assert_has_c10d_ops(`
`132`	`132`	`tp_exported_program.graph_module,`
`133`	`133`	`{`
`134`		`- "_c10d_functional.all_reduce.default": 2,`
`135`		`- "_c10d_functional.wait_tensor.default": 2,`
	`134`	`+ "c10d_functional.all_reduce.default": 2,`
	`135`	`+ "c10d_functional.wait_tensor.default": 2,`
`136`	`136`	`},`
`137`	`137`	`)`
`138`	`138`
`@@ -167,8 +167,8 @@ def test_tp_transform_no_bias(self):`
`167`	`167`	`self.assert_has_c10d_ops(`
`168`	`168`	`tp_exported_program.graph_module,`
`169`	`169`	`{`
`170`		`- "_c10d_functional.all_reduce.default": 1,`
`171`		`- "_c10d_functional.wait_tensor.default": 1,`
	`170`	`+ "c10d_functional.all_reduce.default": 1,`
	`171`	`+ "c10d_functional.wait_tensor.default": 1,`
`172`	`172`	`},`
`173`	`173`	`)`
`174`	`174`