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()
