| 49 | @pytest.mark.require_ngpu(2) |
| 50 | @pytest.mark.isolated_distributed |
| 51 | def test_dist_grad(): |
| 52 | world_size = 2 |
| 53 | x_np = np.random.rand(10).astype("float32") |
| 54 | |
| 55 | @dist.launcher |
| 56 | def worker(): |
| 57 | rank = dist.get_rank() |
| 58 | if rank == 0: |
| 59 | with Grad() as grad: |
| 60 | x = as_tensor(x_np) |
| 61 | grad.wrt(x, callback=save_to(x)) |
| 62 | # need a placeholder to trace operator |
| 63 | remote_send(x, 1) |
| 64 | recv_x = remote_recv(1) |
| 65 | y = recv_x * recv_x |
| 66 | grad([y], [as_tensor(np.ones_like(x_np))]) |
| 67 | np.testing.assert_almost_equal(x.grad.numpy(), x.numpy() * 2) |
| 68 | elif rank == 1: |
| 69 | with Grad() as grad: |
| 70 | recv_x = remote_recv(0) |
| 71 | remote_send(recv_x, 0) |
| 72 | grad([], []) |
| 73 | |
| 74 | worker() |
| 75 | |
| 76 | |
| 77 | def test_grad(): |