add unit test for other_ops

KingsleyLiu-NV · KingsleyLiu-NV · commit 7ba32958c609 · 2023-11-19T21:33:27.000-08:00
diff --git a/test/distributed/_tensor/test_other_ops.py b/test/distributed/_tensor/test_other_ops.py
@@ -0,0 +1,196 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+# Owner(s): ["oncall: distributed"]
+
+
+import torch
+
+import torch.distributed as dist
+
+from timm.models.layers import DropPath
+from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate
+from torch.testing._internal.common_utils import run_tests
+from torch.testing._internal.distributed._tensor.common_dtensor import (
+    DTensorTestBase,
+    with_comms,
+)
+
+ITER_TIME = 10
+LR = 0.001
+
+
+class DistOtherOpsTest(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        # hard code world size to 2
+        return 2
+
+    @with_comms
+    def test_slice(self):
+        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
+        shard_spec = [Replicate()]
+
+        input_list = torch.rand(ITER_TIME, 1024, 10)
+        grad_output_list = torch.rand(ITER_TIME, 1024, 5) * 1e-3
+
+        for i in range(ITER_TIME):
+            inp = input_list[i].to(self.device_type).requires_grad_()
+            grad_output = grad_output_list[i].to(self.device_type)
+
+            # droppath  with dtensor
+            inp_dtensor = distribute_tensor(inp, device_mesh, shard_spec)
+            grad_output_dtensor = distribute_tensor(
+                grad_output, device_mesh, shard_spec
+            )
+            output = inp_dtensor[:, :5]
+            output.backward(grad_output_dtensor)
+
+            # nll with plain tensor
+            output_gt = inp[:, :5]
+            output_gt.backward(grad_output)
+
+            output_diff_abs = output.to_local() - output_gt
+            output_diff_rel = output_diff_abs / (torch.abs(output_gt) + 1e-8)
+            output_mse_abs = torch.mean(output_diff_abs * output_diff_abs).item()
+            output_mse_rel = torch.mean(output_diff_rel * output_diff_rel).item()
+
+            grad_diff_abs = inp_dtensor.grad.to_local() - inp.grad
+            grad_diff_rel = grad_diff_abs / (torch.abs(inp.grad) + 1e-8)
+            grad_mse_abs = torch.mean(grad_diff_abs * grad_diff_abs).item()
+            grad_mse_rel = torch.mean(grad_diff_rel * grad_diff_rel).item()
+
+            self.assertTrue(
+                output_mse_abs <= 1e-6,
+                f"Too large absolute mse for output, expected less equal 1e-6, got {output_mse_abs}",
+            )
+            self.assertTrue(
+                output_mse_rel <= 1e-6,
+                f"Too large relative mse for output, expected less equal 1e-6, got {output_mse_rel}",
+            )
+            self.assertTrue(
+                grad_mse_abs <= 1e-6,
+                f"Too large absolute mse for gradient, expected less equal 1e-6, got {grad_mse_abs}",
+            )
+            self.assertTrue(
+                grad_mse_rel <= 1e-6,
+                f"Too large relative mse for gradient, expected less equal 1e-6, got {grad_mse_rel}",
+            )
+
+    @with_comms
+    def test_bernoulli(self):
+        rank = dist.get_rank()
+        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
+        shard_spec = [Replicate()]
+
+        input_list = torch.rand(ITER_TIME, 1024, 10)
+        grad_output_list = torch.rand(ITER_TIME, 1024, 10) * 1e-3
+
+        droppath = DropPath(0.3)
+
+        for i in range(ITER_TIME):
+            inp = input_list[i].to(self.device_type).requires_grad_()
+            grad_output = grad_output_list[i].to(self.device_type)
+
+            # droppath  with dtensor
+            inp_dtensor = distribute_tensor(inp, device_mesh, shard_spec)
+            grad_output_dtensor = distribute_tensor(
+                grad_output, device_mesh, shard_spec
+            )
+            output = droppath(inp_dtensor)
+            output.backward(grad_output_dtensor)
+
+            send_output_tensor = output.to_local()
+            recv_output_tensor = torch.zeros_like(send_output_tensor)
+
+            send_grad_tensor = inp_dtensor.grad.to_local()
+            recv_grad_tensor = torch.zeros_like(send_grad_tensor)
+
+            send_op_1 = dist.P2POp(dist.isend, send_output_tensor, 1 ^ rank)
+            send_op_2 = dist.P2POp(dist.isend, send_grad_tensor, 1 ^ rank)
+            recv_op_1 = dist.P2POp(dist.irecv, recv_output_tensor, 1 ^ rank)
+            recv_op_2 = dist.P2POp(dist.irecv, recv_grad_tensor, 1 ^ rank)
+
+            reqs = dist.batch_isend_irecv([send_op_1, send_op_2, recv_op_1, recv_op_2])
+            for req in reqs:
+                req.wait()
+
+            output_diff_abs = send_output_tensor - recv_output_tensor
+            output_diff_rel = output_diff_abs / (torch.abs(recv_output_tensor) + 1e-8)
+            output_mse_abs = torch.mean(output_diff_abs * output_diff_abs).item()
+            output_mse_rel = torch.mean(output_diff_rel * output_diff_rel).item()
+
+            grad_diff_abs = send_grad_tensor - recv_grad_tensor
+            grad_diff_rel = grad_diff_abs / (torch.abs(recv_grad_tensor) + 1e-8)
+            grad_mse_abs = torch.mean(grad_diff_abs * grad_diff_abs).item()
+            grad_mse_rel = torch.mean(grad_diff_rel * grad_diff_rel).item()
+
+            self.assertTrue(
+                output_mse_abs <= 1e-6,
+                f"Too large absolute mse for output, expected less equal 1e-6, got {output_mse_abs}",
+            )
+            self.assertTrue(
+                output_mse_rel <= 1e-6,
+                f"Too large relative mse for output, expected less equal 1e-6, got {output_mse_rel}",
+            )
+            self.assertTrue(
+                grad_mse_abs <= 1e-6,
+                f"Too large absolute mse for gradient, expected less equal 1e-6, got {grad_mse_abs}",
+            )
+            self.assertTrue(
+                grad_mse_rel <= 1e-6,
+                f"Too large relative mse for gradient, expected less equal 1e-6, got {grad_mse_rel}",
+            )
+
+    @with_comms
+    def test_nll(self):
+        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
+        shard_spec = [Replicate()]
+
+        pred_list = torch.rand(ITER_TIME, 1024, 10)
+        target_list = torch.randint(0, 10, (ITER_TIME, 1024), dtype=torch.long)
+
+        criterion = torch.nn.CrossEntropyLoss()
+
+        for i in range(ITER_TIME):
+            pred = pred_list[i].to(self.device_type).requires_grad_()
+            target = target_list[i].to(self.device_type)
+
+            # nll with dtensor
+            pred_dtensor = distribute_tensor(pred, device_mesh, shard_spec)
+            target_dtensor = distribute_tensor(target, device_mesh, shard_spec)
+            loss = criterion(pred_dtensor, target_dtensor)
+            loss.backward()
+
+            # nll with plain tensor
+            loss_gt = criterion(pred, target)
+            loss_gt.backward()
+
+            loss_diff_abs = loss.to_local() - loss_gt
+            loss_diff_rel = loss_diff_abs / (torch.abs(loss_gt) + 1e-8)
+            loss_mse_abs = torch.mean(loss_diff_abs * loss_diff_abs).item()
+            loss_mse_rel = torch.mean(loss_diff_rel * loss_diff_rel).item()
+
+            grad_diff_abs = pred_dtensor.grad.to_local() - pred.grad
+            grad_diff_rel = grad_diff_abs / (torch.abs(pred.grad) + 1e-8)
+            grad_mse_abs = torch.mean(grad_diff_abs * grad_diff_abs).item()
+            grad_mse_rel = torch.mean(grad_diff_rel * grad_diff_rel).item()
+
+            self.assertTrue(
+                loss_mse_abs <= 1e-6,
+                f"Too large absolute mse for loss, expected less equal 1e-6, got {loss_mse_abs}",
+            )
+            self.assertTrue(
+                loss_mse_rel <= 1e-6,
+                f"Too large relative mse for loss, expected less equal 1e-6, got {loss_mse_rel}",
+            )
+            self.assertTrue(
+                grad_mse_abs <= 1e-6,
+                f"Too large absolute mse for gradient, expected less equal 1e-6, got {grad_mse_abs}",
+            )
+            self.assertTrue(
+                grad_mse_rel <= 1e-6,
+                f"Too large relative mse for gradient, expected less equal 1e-6, got {grad_mse_rel}",
+            )
+
+
+if __name__ == "__main__":
+    run_tests()
