I looked at the code and did not see something suspicious.
However, the data corruption is still there.
I found that if I do torch.distributed.synchronize(device) explicitly before Isends the problem is mitigated and can be mistaken to “solved”, but I don’t like this solution at all.
I don’t see any rational reason to do so, there is probably some bug there.
Reading the warnings here and here makes me believe that the MPI/distributed API probably does not do many stuff necessary for sharing tensors
like handling references counts, using mutex to guard stuff and etc.
I use CUDA-aware with openMPI, I thought it is supported.