[ca] add test_dtensor_compile.py to compiled autograd tests (pytorch#144107)

more than half the tests use autograd, pass rate 19/26

Pull Request resolved: pytorch#144107
Approved by: https://github.com/zou3519, https://github.com/bdhirsh, https://github.com/jansel

Author: xmfan

test/distributed/_tensor/test_dtensor_compile.py

@@ -86,14 +86,18 @@ def extract_graph(fx_g, _, graph_cell):
 
 class TestDTensorCompile(torch._dynamo.test_case.TestCase):
     def setUp(self):
-        super().setUp()
+        super(
+            type(self), self
+        ).setUp()  # use explicit params for compiled autograd test wrapping
         fake_store = FakeStore()
         dist.init_process_group(
             "fake", store=fake_store, rank=0, world_size=self.world_size
         )
 
     def tearDown(self):
-        super().tearDown()
+        super(
+            type(self), self
+        ).tearDown()  # use explicit params for compiled autograd test wrapping
         dist.destroy_process_group()
 
     @property
@@ -104,6 +108,19 @@ def device_type(self) -> str:
     def world_size(self) -> int:
         return 2
 
+    def test_dtensor_basic(self):
+        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+
+        @torch.compile(backend="aot_eager", fullgraph=True)
+        def fn(x):
+            return x * x + 2
+
+        param = torch.randn(4, 4, requires_grad=True)
+        x = DTensor.from_local(param, mesh, [Shard(0)], run_check=False)
+
+        res = fn(x)
+        res.to_local().sum().backward()
+
     def test_placement_compile(self):
         def fn(x):
             a = 0

test/inductor/test_compiled_autograd.py

@@ -3594,6 +3594,14 @@ def wrap_test_class(orig_cls):
     "test_invalid_gradients",  # can't give autograd error due to inaccurate output metadata of lifted backward
     "test_autograd_node_isinstance",  # backward ctx is a fake cls and not directly a Node instance
     "test_backward_hook_relative_ordering",  # compiled autograd collects breadth first, and module backward hook not supported
+    # Category: Subclasses
+    "test_dtensor_basic",
+    "test_dtensor_contiguous_dtensor_noncontiguous_local_as_tangent",
+    "test_dtensor_different_gradient_placement",
+    "test_dtensor_noncontiguous_output",
+    "test_dtensor_partial_placement_graph_output",
+    "test_tp_compile_comm_reordering",
+    "test_unwrap_async_collective_tensor_tangent",
     # Uncategorized
 }
 
@@ -3606,6 +3614,11 @@ def wrap_test_class(orig_cls):
 
 TestAutogradWithCompiledAutograd = wrap_test_class(test_autograd.TestAutograd)
 TestCustomOpWithCompiledAutograd = wrap_test_class(test_custom_ops.TestCustomOp)
+if torch.distributed.is_available() and HAS_CUDA:
+    test_dtensor = load_test_module("distributed/_tensor/test_dtensor_compile")
+    TestDTensorCompileWithCompiledAutograd = wrap_test_class(
+        test_dtensor.TestDTensorCompile
+    )
 
 if __name__ == "__main__":
     if HAS_CPU: