Update on "[dtensor] make replicate -> partial do division instead"

wanchaol · wanchaol · commit 75a68e2f55ac · 2023-10-09T15:35:28.000-07:00
This PR switches the replicate -&gt; partial to do division instead of
zeroing out other ranks, it preserve same numerics, but avoid the
per-rank behavior difference, and friendly to torch compile

[ghstack-poisoned]
diff --git a/test/distributed/_tensor/test_redistribute.py b/test/distributed/_tensor/test_redistribute.py
@@ -116,7 +116,7 @@ def test_partial_to_replicate_forward_backward(self):
         # replicate to partial internally, and also partial to replicate
         # backward should work as expected
         device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
-        partial_local = torch.randn(12, 3, device=self.device_type, requires_grad=True)
+        partial_local = torch.ones(12, 3, device=self.device_type, requires_grad=True)
         partial_spec = [_Partial()]
         replica_spec = [Replicate()]
         # test partial -> replicate, which trigger all_reduce
@@ -131,8 +131,9 @@ def test_partial_to_replicate_forward_backward(self):
         # test backward to have replicate grad on partial
         global_partial_tensor.backward(torch.ones_like(global_partial_tensor))
         self.assertIsNotNone(partial_local.grad)
-        if device_mesh.get_rank() == 0:
-            self.assertEqual(partial_local.grad, torch.ones_like(partial_local))
+        self.assertEqual(
+            partial_local.grad, torch.ones_like(partial_local) / self.world_size
+        )
 
     @with_comms
     def test_replicate_to_partial(self):
