Skip to content

Instantly share code, notes, and snippets.

@dasayan05
Last active December 4, 2019 17:45
Show Gist options
  • Save dasayan05/a295ef736e1e928ec1cabe046cf13f4e to your computer and use it in GitHub Desktop.
Save dasayan05/a295ef736e1e928ec1cabe046cf13f4e to your computer and use it in GitHub Desktop.
Peer-to-peer communication
# filename 'ptdist.py'
import torch
import torch.distributed as dist
def main(rank, world):
if rank == 0:
x = torch.tensor([1., -1.]) # Tensor of interest
dist.send(x, dst=1)
print('Rank-0 has sent the following tensor to Rank-1')
print(x)
else:
z = torch.tensor([0., 0.]) # A holder for recieving the tensor
dist.recv(z, src=0)
print('Rank-1 has recieved the following tensor from Rank-0')
print(z)
if __name__ == '__main__':
dist.init_process_group(backend='mpi')
main(dist.get_rank(), dist.get_world_size())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment