RocketFlash:

could you give some example code how to use sharing data between ranks, please?


One option is to use the collective communication APIs, e.g., all_gather 25 (NCCL and Gloo) or gather 3 (Gloo). Sth like:
# on all ranks

out_tensors = [torch.zeros(2, 2), torch.zeros(2, 2)]
inp_tensor = [torch.ones(2, 2)]

torch.distributed.all_gather(out_tensors, inp_tensor)