Avoid nested CommTensor wrapping

mrshenli · mrshenli · commit 61e89dbfff41 · 2022-09-13T22:59:13.000Z
[ghstack-poisoned]
diff --git a/test/distributed/test_c10d_common.py b/test/distributed/test_c10d_common.py
@@ -1472,6 +1472,13 @@ def comm_fn(tensor, group=None):
 
         self._test_work_wait(tensor, comm_fn=comm_fn)
 
+    def _test_nested_comm_tensor_wrapping(self, tensor):
+        def comm_fn(tensor, group=None):
+            work = dist.all_reduce(CommTensor(tensor), group=group, async_op=True)
+            return work, tensor
+
+        self._test_work_wait(tensor, comm_fn=comm_fn)
+
 
 if __name__ == "__main__":
     assert (
diff --git a/test/distributed/test_c10d_gloo.py b/test/distributed/test_c10d_gloo.py
@@ -2396,6 +2396,9 @@ def test_scatter_work_wait_gpu(self):
             torch.ones(2, 2, device=self.rank) * self.rank
         )
 
+    def test_nested_comm_tensor_wrapping(self):
+        self._test_nested_comm_tensor_wrapping(torch.ones(2, 2) * self.rank)
+
 
 if __name__ == "__main__":
     assert (
diff --git a/test/distributed/test_c10d_nccl.py b/test/distributed/test_c10d_nccl.py
@@ -2853,6 +2853,11 @@ def test_scatter_work_wait_gpu(self):
             torch.ones(2, 2, device=self.rank) * self.rank
         )
 
+    @skip_if_lt_x_gpu(2)
+    def test_nested_comm_tensor_wrapping(self):
+        self._test_nested_comm_tensor_wrapping(
+            torch.ones(2, 2, device=self.rank) * self.rank
+        )
 
 if __name__ == "__main__":
     assert (
diff --git a/torch/distributed/_spmd/comm_tensor.py b/torch/distributed/_spmd/comm_tensor.py
@@ -93,7 +93,8 @@ class CommTensor(torch.Tensor):
     def __new__(cls, tensor: torch.Tensor):
         r = torch.Tensor._make_subclass(
             cls,
-            tensor,
+            # avoid nested CommTensor Wrapping
+            tensor._tensor if isinstance(tensor, CommTensor) else tensor,
             require_grad=tensor.requires_grad,
         )
         # The tensor object wrapped by this CommTensor

Original file line number	Diff line number	Diff line change
`@@ -2396,6 +2396,9 @@ def test_scatter_work_wait_gpu(self):`
`2396`	`2396`	`torch.ones(2, 2, device=self.rank) * self.rank`
`2397`	`2397`	`)`
`2398`	`2398`
	`2399`	`+ def test_nested_comm_tensor_wrapping(self):`
	`2400`	`+ self._test_nested_comm_tensor_wrapping(torch.ones(2, 2) * self.rank)`
	`2401`	`+`
`2399`	`2402`
`2400`	`2403`	`if __name__ == "__main__":`
`2401`	`2404`	`assert (`