Repository: pytorch
Issue: [CUDA RPC] Incorrect results of GPU Tensor transferring using RPC when parallelized with other GPU programs
Labels: oncall: distributed,module: cuda

### 🐛 Describe the bug


### Issue Summary
This is a simplified version of [Issue](https://github.com/pytorch/pytorch/issues/100725#issue-1697886172).

I transfer a GPU tensor between two GPUs using PyTorch RPC, and find incorrect results when parallelized with other GPU programs, such as a pure computation task in the backend using PyTorch. When the other GPU task is stopped, the results become right.

This suggests that other GPU programs could be interfering with the CUDA Support RPC, leading to incorrect message transfers.

Could anyone provide some insight or guidance on how to prevent other GPU programs from interfering with the CUDA Support RPC and ensure correct message transfers? Any help or suggestions would be greatly appreciated.

### Steps tp Reproduce
1. Run a pure computation task in the backend using PyTorch.
```python
import torch
import time
from multiprocessing import Pool, set_start_method


def run_on_single_gpu(device):
    a = torch.randn(20000,20000).cuda(device)
    b = torch.randn(20000,20000).cuda(device)
    ta = a
    tb = b
    while True:
        a = ta
        b = tb
        a = torch.sin(a)
        b = torch.sin(b)
        a = torch.cos(a)
        b = torch.cos(b)
        a = torch.tan(a)
        b = torch.tan(b)
        a = torch.exp(a)
        b = torch.exp(b)
        a = torch.log(a)
        b = torch.log(b)
        b = torch.matmul(a, b)
        #time.sleep(0.000005)

if __name__ == '__main__':
    set_start_method('spawn')
    print('start running')
    num_gpus = torch.cuda.device_count()
    pool = Pool(processes=num_gpus)
    pool.map(run_on_single_gpu, range(num_gpus))
    pool.close()
    pool.join()
   ```
2. Run a PyTorch program which  transfers GPU tensors using RPC.
```python
import os
import time
from functools import wraps
import torch
import torch.nn as nn
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp
from torch.distributed.rpc import RRef

class Identity(nn.Module):
    def __init__(self, device):
        super(Identity, self).__init__()
        self.device = device

    def forward(self, x_rref):
        out = x_rref.to_here().to(self.device)
        # return out.cpu()
        return out



class DistP2P(nn.Module):
    """
    Assemble two parts as an nn.Module and define pipelining logic
    """
    def __init__(self, workers):
        super(DistP2P, self).__init__()

        # Put the first part of the ResNet50 on workers[0]
        self.p1_rref = rpc.remote(
            workers[0],
            Identity,
            args = ("cuda:0",)
        )

        # Put the second part of the ResNet50 on workers[1]
        self.p2_rref = rpc.remote(
            workers[1],
            Identity,
            args = ("cuda:1",)
        )

    def forward(self, xs):
        out_futures = []
        x_rref = RRef(xs)
        y_rref = self.p1_rref.remote().forward(x_rref)
        z_fut = self.p2_rref.rpc_async().forward(y_rref)
        out_futures.append(z_fut)

        # collect and cat all output tensors into one tensor.
        return torch.cat(torch.futures.wait_all(out_futures))



#########################################################
#                   Run RPC Processes                   #
#########################################################

def run_master():

    # put the two model parts on worker1 and worker2 respectively
    p2p = DistP2P(["worker1", "worker2"])
    torch.manual_seed(1234)
    inputs = torch.randn(1, 12)

    for i in range(5):
        print(f"Processing batch {i}")
        outputs = p2p(inputs)
        print(f"outputs: {outputs}")


def run_worker(rank, world_size):
    os.environ['MASTER_ADDR'] = '11.218.124.179'
    os.environ['MASTER_PORT'] = '29000'

    # Higher timeout is added to accommodate for kernel compilation time in case of ROCm.
    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256, rpc_timeout=300)

    if rank == 0:
        options.set_device_map("master", {rank: 0})
        options.set_device_map("worker1", {rank: 0})
        options.set_device_map("worker2", {rank: 1})
        rpc.init_rpc(
            "master",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=options
        )
        run_master()
    else:
        options.set_device_map("master", {rank-1: 0})
        options.set_device_map("worker1", {rank-1: 0})
        options.set_device_map("worker2", {rank-1: 1})
        rpc.init_rpc(
            f"worker{rank}",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=options
        )
        pass

    # block until all rpcs finish
    rpc.shutdown()


if __name__=="__main__":
    world_size = 3
    mp.spawn(run_worker, args=(world_size,), nprocs=world_size, join=True)
```
3. Observe the log of the RPC task and note that the output values are inconsistent when I fix the transferred GPU tensor.
<img width="801" alt="image" src="https://user-images.githubusercontent.com/34735292/236594605-aa5ae209-4bbe-4f57-99dc-e7a70c584897.png">

4. Close the pure computation task and only run the RPC task. Observe that the output values become consistent.  
<img width="835" alt="image" src="https://user-images.githubusercontent.com/34735292/236594630-ab874480-de37-4e49-b552-1c40abd2d4cc.png">


### Versions

PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Tencent tlinux 2.2 (Final) (x86_64)
GCC version: (GCC) 8.3.0
Clang version: Could not collect
CMake version: version 3.26.1
Libc version: glibc-2.17

Python version: 3.8.12 (default, Jun 13 2022, 19:37:57)  [GCC 8.3.0] (64-bit runtime)
Python platform: Linux-4.14.105-1-tlinux3-0013-x86_64-with-glibc2.2.5
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: A100-SXM4-40GB
GPU 1: A100-SXM4-40GB
GPU 2: A100-SXM4-40GB
GPU 3: A100-SXM4-40GB
GPU 4: A100-SXM4-40GB
GPU 5: A100-SXM4-40GB
GPU 6: A100-SXM4-40GB
GPU 7: A100-SXM4-40GB

Nvidia driver version: 450.156.00
cuDNN version: Probably one of the following:
/usr/lib64/libcudnn.so.8.5.0
/usr/lib64/libcudnn_adv_infer.so.8.5.0
/usr/lib64/libcudnn_adv_train.so.8.5.0
/usr/lib64/libcudnn_cnn_infer.so.8.5.0
/usr/lib64/libcudnn_cnn_train.so.8.5.0
/usr/lib64/libcudnn_ops_infer.so.8.5.0
/usr/lib64/libcudnn_ops_train.so.8.5.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:          x86_64
CPU op-mode(s):        32-bit, 64-bit
Byte Order:            Little Endian
CPU(s):                192
On-line CPU(s) list:   0-191
Thread(s) per core:    2
Core(s) per socket:    48
Socket(s):             2
NUMA node(s):          2
Vendor ID:             AuthenticAMD
CPU family:            23
Model:                 49
Model name:            AMD EPYC 7K62 48-Core Processor
Stepping:              0
CPU MHz:               3301.340
CPU max MHz:           2600.0000
CPU min MHz:           1500.0000
BogoMIPS:              5190.52
Virtualization:        AMD-V
L1d cache:             32K
L1i cache:             32K
L2 cache:              512K
L3 cache:              16384K
NUMA node0 CPU(s):     0-47,96-143
NUMA node1 CPU(s):     48-95,144-191
Flags:                 fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif umip rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.0.0
[pip3] torchaudio==2.0.1
[pip3] torchpippy==0.1.0
[pip3] torchvision==0.12.0
[pip3] triton==2.0.0
[conda] Could not collect

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @ngimel