Repository: pytorch
Issue: [CUDA RPC] Incorrect messages in CUDA Support RPC when parallelized with other GPU programs
Labels: oncall: distributed

### 🐛 Describe the bug

### Issue Summary
When running the official example of [DISTRIBUTED PIPELINE PARALLELISM USING RPC](https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html) and modifying the transferred tensors into GPU tensors (i.e., return out.cpu() -> return out), inconsistencies in the transferred tensor values are observed. Specifically, there are discrepancies in the loss values for the same inputs and model parameters.

These discrepancies only occur when running other GPU programs, such as a pure computation task in the backend using PyTorch or nccl-tests tasks. When the other GPU task is stopped, the loss values become consistent. 

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 to 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 the official example of [DISTRIBUTED PIPELINE PARALLELISM USING RPC](https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html) and modifying the transferred tensors into GPU tensors (i.e., return out.cpu() -> return out),

```python
import os
import threading
import time
from functools import wraps

import torch
import torch.nn as nn
import torch.distributed.autograd as dist_autograd
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp
import torch.optim as optim
from torch.distributed.optim import DistributedOptimizer
from torch.distributed.rpc import RRef

from torchvision.models.resnet import Bottleneck


#########################################################
#           Define Model Parallel ResNet50              #
#########################################################

# In order to split the ResNet50 and place it on two different workers, we
# implement it in two model shards. The ResNetBase class defines common
# attributes and methods shared by two shards. ResNetShard1 and ResNetShard2
# contain two partitions of the model layers respectively.


num_classes = 1000


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

class ResNetBase(nn.Module):
    def __init__(self, block, inplanes, num_classes=1000,
                 groups=1, width_per_group=64, norm_layer=None):
        super(ResNetBase, self).__init__()

        self._lock = threading.Lock()
        self._block = block
        self._norm_layer = nn.BatchNorm2d
        self.inplanes = inplanes
        self.dilation = 1
        self.groups = groups
        self.base_width = width_per_group

    def _make_layer(self, planes, blocks, stride=1):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if stride != 1 or self.inplanes != planes * self._block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * self._block.expansion, stride),
                norm_layer(planes * self._block.expansion),
            )

        layers = []
        layers.append(self._block(self.inplanes, planes, stride, downsample, self.groups,
                                  self.base_width, previous_dilation, norm_layer))
        self.inplanes = planes * self._block.expansion
        for _ in range(1, blocks):
            layers.append(self._block(self.inplanes, planes, groups=self.groups,
                                      base_width=self.base_width, dilation=self.dilation,
                                      norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def parameter_rrefs(self):
        r"""
        Create one RRef for each parameter in the given local module, and return a
        list of RRefs.
        """
        return [RRef(p) for p in self.parameters()]


class ResNetShard1(ResNetBase):
    """
    The first part of ResNet.
    """
    def __init__(self, device, *args, **kwargs):
        super(ResNetShard1, self).__init__(
            Bottleneck, 64, num_classes=num_classes, *args, **kwargs)

        self.device = device
        self.seq = nn.Sequential(
            nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False),
            self._norm_layer(self.inplanes),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            self._make_layer(64, 3),
            self._make_layer(128, 4, stride=2)
        ).to(self.device)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x_rref):
        x = x_rref.to_here().to(self.device)
        with self._lock:
            out =  self.seq(x)
        # return out.cpu()
        return out


class ResNetShard2(ResNetBase):
    """
    The second part of ResNet.
    """
    def __init__(self, device, *args, **kwargs):
        super(ResNetShard2, self).__init__(
            Bottleneck, 512, num_classes=num_classes, *args, **kwargs)

        self.device = device
        self.seq = nn.Sequential(
            self._make_layer(256, 6, stride=2),
            self._make_layer(512, 3, stride=2),
            nn.AdaptiveAvgPool2d((1, 1)),
        ).to(self.device)

        self.fc =  nn.Linear(512 * self._block.expansion, num_classes).to(self.device)

    def forward(self, x_rref):
        x = x_rref.to_here().to(self.device)
        with self._lock:
            out = self.fc(torch.flatten(self.seq(x), 1))
        # return out.cpu()
        return out


class DistResNet50(nn.Module):
    """
    Assemble two parts as an nn.Module and define pipelining logic
    """
    def __init__(self, split_size, workers, *args, **kwargs):
        super(DistResNet50, self).__init__()

        self.split_size = split_size

        # Put the first part of the ResNet50 on workers[0]
        self.p1_rref = rpc.remote(
            workers[0],
            ResNetShard1,
            args = ("cuda:0",) + args,
            kwargs = kwargs
        )

        # Put the second part of the ResNet50 on workers[1]
        self.p2_rref = rpc.remote(
            workers[1],
            ResNetShard2,
            args = ("cuda:1",) + args,
            kwargs = kwargs
        )

    def forward(self, xs):
        # Split the input batch xs into micro-batches, and collect async RPC
        # futures into a list
        out_futures = []
        for x in iter(xs.split(self.split_size, dim=0)):
            x_rref = RRef(x)
            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))

    def parameter_rrefs(self):
        remote_params = []
        remote_params.extend(self.p1_rref.remote().parameter_rrefs().to_here())
        remote_params.extend(self.p2_rref.remote().parameter_rrefs().to_here())
        return remote_params


#########################################################
#                   Run RPC Processes                   #
#########################################################

num_batches = 3
batch_size = 120
image_w = 128
image_h = 128


def run_master(split_size):

    # put the two model parts on worker1 and worker2 respectively
    model = DistResNet50(split_size, ["worker1", "worker2"])
    loss_fn = nn.MSELoss()
    opt = DistributedOptimizer(
        optim.SGD,
        model.parameter_rrefs(),
        lr=0.05,
    )

    one_hot_indices = torch.LongTensor(batch_size) \
                           .random_(0, num_classes) \
                           .view(batch_size, 1)

    for i in range(num_batches):
        print(f"Processing batch {i}")
        # generate random inputs and labels
        torch.manual_seed(123)
        inputs = torch.randn(batch_size, 3, image_w, image_h)
        labels = torch.zeros(batch_size, num_classes) \
                      .scatter_(1, one_hot_indices, 1)

        # The distributed autograd context is the dedicated scope for the
        # distributed backward pass to store gradients, which can later be
        # retrieved using the context_id by the distributed optimizer.
        with dist_autograd.context() as context_id:
            outputs = model(inputs)
            loss = loss_fn(outputs, labels.to("cuda:0"))
            print(f"loss: {loss.item()}")
            dist_autograd.backward(context_id, [loss])
            # opt.step(context_id)


def run_worker(rank, world_size, num_split):
    os.environ['MASTER_ADDR'] = '11.218.124.179'
    # os.environ['MASTER_ADDR'] = 'localhost'
    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(num_split)
    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
    for num_split in [1, 2, 4, 8]:
        tik = time.time()
        mp.spawn(run_worker, args=(world_size, num_split), nprocs=world_size, join=True)
        tok = time.time()
        print(f"number of splits = {num_split}, execution time = {tok - tik}")
```

3. Observe the log of the RPC task and note that the loss values are inconsistent. For example, the losses in 'number of splits = 4 and 8' of my log.

```Processing batch 0
loss: 0.3522467613220215
Processing batch 1
loss: 0.3522467613220215
Processing batch 2
loss: 0.3522467613220215
number of splits = 1, execution time = 30.63558030128479
Processing batch 0
loss: 0.37125730514526367
Processing batch 1
loss: 0.37125730514526367
Processing batch 2
loss: 0.37125730514526367
number of splits = 2, execution time = 25.16931390762329
Processing batch 0
loss: 0.3850472569465637
Processing batch 1
loss: 0.3850732445716858
Processing batch 2
loss: 0.38506922125816345
number of splits = 4, execution time = 24.290683269500732
Processing batch 0
loss: 0.38254597783088684
Processing batch 1
loss: 0.38254597783088684
Processing batch 2
loss: 0.35708141326904297
number of splits = 8, execution time = 22.055739402770996
```

4. Close the pure computation task and only run the RPC task. Observe that the loss values become consistent.
```
Processing batch 0
loss: 0.3423759937286377
Processing batch 1
loss: 0.3423759937286377
Processing batch 2
loss: 0.3423759937286377
number of splits = 1, execution time = 19.911068439483643
Processing batch 0
loss: 0.40370798110961914
Processing batch 1
loss: 0.40370798110961914
Processing batch 2
loss: 0.40370798110961914
number of splits = 2, execution time = 14.234280824661255
Processing batch 0
loss: 0.36699122190475464
Processing batch 1
loss: 0.36699122190475464
Processing batch 2
loss: 0.36699122190475464
number of splits = 4, execution time = 12.719268560409546
Processing batch 0
loss: 0.3601795434951782
Processing batch 1
loss: 0.3601795434951782
Processing batch 2
loss: 0.3601795434951782
number of splits = 8, execution time = 11.449119329452515
```

### 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:               3292.824
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