def exchange(model, world_size, rank): old_model = copy.deepcopy(model) for param in old_model.parameters(): dist.isend( param.data, dst=(rank+1)%world_size ) for param in model.parameters(): dist.recv( param.data, src=(rank-1)%world_size ) return model
def test_isend(self): rank = dist.get_rank() world_size = dist.get_world_size() if rank == 0: requests = [ dist.isend(_build_tensor(dest, 10), dest) for dest in range(1, world_size) ] for request in requests: request.wait() self.assertTrue(request.is_completed()) else: tensor = _build_tensor(rank, -1) dist.recv(tensor, 0) self.assertEqual(tensor, _build_tensor(rank, 10)) self._barrier()