add unittest for dpp gather autograd compatibility

tests/unittests/bases/test_ddp.py

@@ -77,6 +77,40 @@ def _test_ddp_gather_uneven_tensors_multidim(rank: int, worldsize: int = NUM_PRO
         assert (val == torch.ones_like(val)).all()
 
 
+def _test_ddp_gather_autograd_same_shape(rank: int, worldsize: int = NUM_PROCESSES) -> None:
+    tensor = torch.ones(50, requires_grad=True)
+    result = gather_all_tensors(tensor)
+    assert len(result) == worldsize
+    scalar1 = 0
+    scalar2 = 0
+    for idx in range(worldsize):
+        if idx == rank:
+            scalar1 = scalar1 + torch.sum(tensor * torch.ones_like(tensor))
+        else:
+            scalar1 = scalar1 + torch.sum(result[idx] * torch.ones_like(result[idx]))
+        scalar2 = scalar2 + torch.sum(result[idx] * torch.ones_like(result[idx]))
+    gradient1 = torch.autograd.grad(scalar1, [tensor], retain_graph=True)[0]
+    gradient2 = torch.autograd.grad(scalar2, [tensor])[0]
+    assert torch.allclose(gradient1, gradient2)
+
+
+def _test_ddp_gather_autograd_different_shape(rank: int, worldsize: int = NUM_PROCESSES) -> None:
+    tensor = torch.ones(rank + 1, 2 - rank, requires_grad=True)
+    result = gather_all_tensors(tensor)
+    assert len(result) == worldsize
+    scalar1 = 0
+    scalar2 = 0
+    for idx in range(worldsize):
+        if idx == rank:
+            scalar1 = scalar1 + torch.sum(tensor * torch.ones_like(tensor))
+        else:
+            scalar1 = scalar1 + torch.sum(result[idx] * torch.ones_like(result[idx]))
+        scalar2 = scalar2 + torch.sum(result[idx] * torch.ones_like(result[idx]))
+    gradient1 = torch.autograd.grad(scalar1, [tensor], retain_graph=True)[0]
+    gradient2 = torch.autograd.grad(scalar2, [tensor])[0]
+    assert torch.allclose(gradient1, gradient2)
+
+
 def _test_ddp_compositional_tensor(rank: int, worldsize: int = NUM_PROCESSES) -> None:
     dummy = DummyMetricSum()
     dummy._reductions = {"x": torch.sum}
@@ -97,6 +131,8 @@ def _test_ddp_compositional_tensor(rank: int, worldsize: int = NUM_PROCESSES) ->
         _test_ddp_sum_cat,
         _test_ddp_gather_uneven_tensors,
         _test_ddp_gather_uneven_tensors_multidim,
+        _test_ddp_gather_autograd_same_shape,
+        _test_ddp_gather_autograd_different_shape,
         _test_ddp_compositional_tensor,
     ],
 )