Dear theevann,
I am trying to implement inter-GPU communiation by using pytorch+mpi+gpu. I have run your code; it runs successfully.
In addition, I have modified your code a little to make that process0 runs on GPU0 and process1 runs on GPU1.
However, I find that the modified code can not run normally because of “Segmentation Fault”. Do you know why?
Following are the modified code.
import os
import socket
import torch
import torch.distributed as dist
from torch.multiprocessing import Process
import platform
def run(rank, size):
if rank == 0:
    tensor = torch.zeros(1).cuda(0)
    # Send the tensor to process 1
    tensor += 1
    dist.send(tensor=tensor, dst=1)
else:
    tensor = torch.zeros(1).cuda(1)
    # Receive tensor from process 0
    dist.recv(tensor=tensor, src=0)
print('Rank ', rank, ' has data ', tensor[0])

def init_processes(fn):
“”" Initialize the distributed environment. “”"
dist.init_process_group(‘mpi’)
rank = dist.get_rank()
size = dist.get_world_size()
print('I am rank ', rank, ’ on ', platform.node())
fn(rank, size)
if name == “main”:
init_processes(run)
Your help will be appreciated. Thank you.