Repository: pytorch
Issue: Synchronization problem in torch.distributed with MPI on CUDA 
Labels: oncall: distributed,triaged

## 🐛 Bug

Using the `send` and `recv` primitives of `torch.distributed` with MPI on CUDA, I sometimes receive different messages than I sent. I also sometimes get `illegal memory access` crashes.

In the example below, it seems that the `torch.zeros_like` is still running while the receive starts, leading to a decrease in mean value of the tensor sent/received.
Adding `torch.cuda.synchronize()` before `dist.recv` fixes the issue.

This only happens when multiple processes run on the same GPU. I expect this to be because `send` / `recv`/ are very fast in this setting (using CUDA shared memory.)

I use OpenMPI 4.0.3rc4 with UCX (see environment below).

## To Reproduce

This snippet reproduces the issue on my machine:

```python
import torch
import torch.distributed as dist

dist.init_process_group("mpi")

rank = dist.get_rank()
world_size = dist.get_world_size()  # occurs from 2 workers onwards
right = (rank + 1) % world_size
left = (rank - 1) % world_size

x = torch.ones(100000, device="cuda").float()  # large size is important, bug on happens CUDA

for i in range(100):
    h = dist.isend(x, left)
    recv_buffer = torch.zeros_like(x)
    # torch.cuda.synchronize() <- this fixes the issue
    dist.recv(recv_buffer, right)
    x[:] = recv_buffer
    h.wait()

print(torch.mean(x).item())  # prints 0.0 instead of 1.0
```

I run this with `mpirun -n 2 --hostfile hostfile python example.py`.


## Expected behavior

The script should print `1.0`, but prints `0.0` instead.


## Environment

PyTorch version: 1.4.0a0+7404463
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.10.2

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 10.2.89
GPU models and configuration: GPU 0: Tesla K80
Nvidia driver version: 440.64.00
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip] msgpack-numpy==0.4.4.3
[pip] numpy==1.18.1
[pip] numpydoc==0.9.2
[pip] torch==1.4.0a0+7404463
[pip] torchtext==0.6.0
[pip] torchvision==0.5.0
[conda] blas                      1.0                         mkl  
[conda] mkl                       2020.0                      166  
[conda] mkl-include               2020.0                      166  
[conda] mkl-service               2.3.0            py37he904b0f_0  
[conda] mkl_fft                   1.0.15           py37ha843d7b_0  
[conda] mkl_random                1.1.0            py37hd6b4f25_0  
[conda] numpy                     1.18.1           py37h4f9e942_0  
[conda] numpy-base                1.18.1           py37hde5b4d6_1  
[conda] numpydoc                  0.9.2                      py_0


## Additional context

## UCX

https://github.com/openucx/ucx.git#c9efc71973c459da426000974a7824bac1fe56e9

```
$ ucx_info -v
# UCT version=1.9.0 revision c9efc71
# configured with: --prefix=/usr/cluster --with-cuda=/usr/local/cuda
```

## OpenMPI

https://github.com/open-mpi/ompi.git#8b4a8cd34cb33f3eb89c22adc46243e9df19d7b0

PyTorch built with:
  - GCC 7.5
  - Intel(R) Math Kernel Library Version 2020.0.0 Product Build 20191122 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v0.21.1 (Git Hash 7d2fd500bc78936d1d648ca713b901012f470dbc)
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - NNPACK is enabled
  - CUDA Runtime 10.2
  - NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_70,code=compute_70
  - CuDNN 7.6.5  (built against CUDA 10.1)
  - Build settings: BLAS=MKL, BUILD_NAMEDTENSOR=OFF, BUILD_TYPE=Release, CXX_FLAGS= -Wno-deprecated -fvisibility-inlines-hidden -fopenmp -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -O2 -fPIC -Wno-narrowing -Wall -Wextra -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-sign-compare -Wno-unused-parameter -Wno-unused-variable -Wno-unused-function -Wno-unused-result -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-stringop-overflow -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Wno-stringop-overflow, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, USE_CUDA=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=ON, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON, USE_STATIC_DISPATCH=OFF, 

cc @ezyang @gchanan @zou3519 @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @xush6528 @osalpekar