pytorch
diff --git a/‎test/distributed/tensor/test_dtensor_compile.py
Lines changed: 7 additions & 24 deletions b/‎test/distributed/tensor/test_dtensor_compile.py
Lines changed: 7 additions & 24 deletions
@@ -1,7 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # Owner(s): ["oncall: distributed"]
 
-import contextlib
 import copy
 import functools
 import unittest
@@ -880,17 +879,9 @@ class TestDTensorCompileE2E(DTensorTestBase):
     def world_size(self):
         return 4
 
-    # multiprocess relies on pickling the source code
-    # so compiled autograd tests can't dynamically wrap this class
-    def _bwd_ctx(self, use_ca):
-        if not use_ca:
-            return contextlib.nullcontext()
-        return torch._dynamo.compiled_autograd._enable(torch.compile)
-
     @with_comms
     @parametrize("is_seq_parallel", [True, False])
-    @parametrize("use_ca", [True, False])
-    def test_tp_compile_fullgraph(self, is_seq_parallel, use_ca):
+    def test_tp_compile_fullgraph(self, is_seq_parallel):
         mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
 
         model = SimpleModel(self.device_type)
@@ -944,15 +935,13 @@ def test_tp_compile_fullgraph(self, is_seq_parallel, use_ca):
         cnt = torch._dynamo.testing.CompileCounterWithBackend("aot_eager")
         compiled_mod = torch.compile(model, backend=cnt, fullgraph=True)
         compiled_out = compiled_mod(inp)
-        with self._bwd_ctx(use_ca):
-            compiled_out.sum().backward()
+        compiled_out.sum().backward()
         self.assertEqual(compiled_out, out)
         self.assertEqual(cnt.frame_count, 1)
 
     @with_comms
     @skip_if_lt_x_gpu(4)
-    @parametrize("use_ca", [True, False])
-    def test_2d_fsdp_tp_compile(self, use_ca):
+    def test_2d_fsdp_tp_compile(self):
         data_parallel_size = 2
         model = SimpleModel(self.device_type)
         model_copy = copy.deepcopy(model)
@@ -995,16 +984,13 @@ def test_2d_fsdp_tp_compile(self, use_ca):
         cnt = torch._dynamo.testing.CompileCounterWithBackend("aot_eager")
         compiled_2d = torch.compile(fsdp_2d, backend=cnt)
         compiled_output = compiled_2d(inp)
-        with self._bwd_ctx(use_ca):
-            compiled_output.sum().backward()
 
         self.assertEqual(out, compiled_output)
         self.assertEqual(cnt.frame_count, 1)
 
     @with_comms
     @skip_if_lt_x_gpu(4)
-    @parametrize("use_ca", [True, False])
-    def test_2d_fsdp_tp_ac_compile(self, use_ca):
+    def test_2d_fsdp_tp_ac_compile(self):
         dp_degree = 2
         tp_degree = self.world_size // dp_degree
         model = SimpleModel(self.device_type)
@@ -1047,17 +1033,15 @@ def test_2d_fsdp_tp_ac_compile(self, use_ca):
 
         # backward pass
         out.sum().backward()
-        with self._bwd_ctx(use_ca):
-            compiled_output.sum().backward()
+        compiled_output.sum().backward()
 
         # compare the gradients:
         for n, p in zip(fsdp_2d.parameters(), compiled_2d.parameters()):
             self.assertEqual(n.grad, p.grad)
 
     @with_comms
     @skip_if_lt_x_gpu(4)
-    @parametrize("use_ca", [True, False])
-    def test_compile_dtensor_redistribute_backward(self, use_ca):
+    def test_compile_dtensor_redistribute_backward(self):
         mesh = DeviceMesh(device_type="cuda", mesh=torch.arange(self.world_size))
 
         def fn(x, y):
@@ -1081,8 +1065,7 @@ def fn(x, y):
 
         # Now run and assert the backward + gradients
         ref.sum().backward()
-        with self._bwd_ctx(use_ca):
-            res.sum().backward()
+        res.sum().backward()
 
         self.assertEqual(x_ref.grad, x.grad)
         self.assertEqual(y_ref.grad, y.grad)