Repository: pytorch
Issue: RPC: wait method of Future object return 0 sometimes in rpc framework
Labels: high priority,oncall: distributed,triaged,module: rpc

### 🐛 Describe the bug

I attempt to compare the performance between `rpc_sync()` and `rpc_async()`  of PyRRef object. The call of `PyRRef.rpc_async()` will return a RRefProxy object and the call of `RRefProxy.forward(x)` will return a torch.Future object. I can call wait method of Future to get the real data. But I found the return value of wait method can sometimes be zero, which is quite confusing. 

The config is:
- world_size=4
- PyTorch: 1.12.0+cu113
- Python: 3.8.10
- backend: 'nccl'
- host: 'localhost'

Codes are below:

```python
def run_worker(rank, host, port, backend):
    init_rpc(rank, host, port)
    
    device = 'cuda'
    if rank == 0:
        x = torch.randn([1000, 1000], device=device)

        # init rref of model of each tsage
        rref_rank0_model: PyRRef = rpc.remote("work0", TestModel, args=(0, 'cpu' if device == 'cpu' else 'cuda:0'))
        rref_rank1_model: PyRRef = rpc.remote("work1", TestModel, args=(1, 'cpu' if device == 'cpu' else 'cuda:1'))
        rref_rank2_model: PyRRef = rpc.remote("work2", TestModel, args=(2, 'cpu' if device == 'cpu' else 'cuda:2'))
        rref_rank3_model: PyRRef = rpc.remote("work3", TestModel, args=(3, 'cpu' if device == 'cpu' else 'cuda:3'))
        
        # init blocking proxy of each rref
        sync_proxy_model_0: RRefProxy = rref_rank0_model.rpc_sync()
        sync_proxy_model_1: RRefProxy = rref_rank1_model.rpc_sync()
        sync_proxy_model_2: RRefProxy = rref_rank2_model.rpc_sync()
        sync_proxy_model_3: RRefProxy = rref_rank3_model.rpc_sync()

        # init non blocking proxy of 
        async_proxy_model_0: RRefProxy = rref_rank0_model.rpc_async()
        async_proxy_model_1: RRefProxy = rref_rank1_model.rpc_async()
        async_proxy_model_2: RRefProxy = rref_rank2_model.rpc_async()
        async_proxy_model_3: RRefProxy = rref_rank3_model.rpc_async()
        
        # run async call 10 times
        for _ in range(10):
            out_0 = async_proxy_model_0.forward(x).wait()
            out_1 = async_proxy_model_1.forward(out_0).wait()
            out_2 = async_proxy_model_2.forward(out_1).wait()
            out_3 = async_proxy_model_3.forward(out_2).wait()
            print(f"async check sum {out_3.sum()}")
            
        # run sync call 10 times
        for _ in range(10):
            out_0: torch.Tensor = sync_proxy_model_0.forward(x)
            out_1: torch.Tensor = sync_proxy_model_1.forward(out_0)        
            out_2: torch.Tensor = sync_proxy_model_2.forward(out_1)
            out_3: torch.Tensor = sync_proxy_model_3.forward(out_2)
            print(f"sync check sum {out_3.sum()}")
        
    # barrier for rpc
    rpc.shutdown()
```

And the outputs are:
```
async check sum 223.30909729003906
async check sum 244.09559631347656
async check sum 244.09559631347656
async check sum 223.30909729003906
async check sum 0.0
async check sum 223.30909729003906
async check sum 251.9927978515625
async check sum 251.9927978515625
async check sum 223.30909729003906
async check sum 251.9927978515625
sync check sum 223.30909729003906
sync check sum 223.30909729003906
sync check sum 223.30909729003906
sync check sum 223.30909729003906
sync check sum 223.30909729003906
sync check sum 223.30909729003906
sync check sum 223.30909729003906
sync check sum 223.30909729003906
sync check sum 223.30909729003906
sync check sum 223.30909729003906
```

The outputs of **check sum** are expected to be all the same, but outputs of rpc_async are different.

The complete code is as follows:

```python
import os

import torch
from torch.distributed import rpc
import torch.multiprocessing as mp
import torch.nn as nn
from torch._C._distributed_rpc import PyRRef
from torch.distributed.rpc.rref_proxy import RRefProxy

world_size = 4
torch.manual_seed(123)

class TestModel(nn.Module):
    def __init__(self, rank, device) -> None:
        super().__init__()
        self._rank = rank
        if rank == 0:
            # self.linear_0 = nn.Linear(1000, 128).to(device)
            self.linear_0 = nn.Sequential(
                nn.Linear(1000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 128)
            ).to(device)
        elif rank == 1:
            # self.linear_1 = nn.Linear(128, 256).to(device)
            self.linear_1 = nn.Sequential(
                nn.Linear(128, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 256),
                nn.Linear(256, 256)
            ).to(device)

        elif rank == 2:
            # self.linear_2 = nn.Linear(256, 16).to(device)
            self.linear_2 = nn.Sequential(
                nn.Linear(256, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 16)
            ).to(device)

        elif rank == 3:
            # self.linear_3 = nn.Linear(16, 1000).to(device)
            self.linear_3 = nn.Sequential(
                nn.Linear(16, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 2000),
                nn.Linear(2000, 1)
            ).to(device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        linear: nn.Linear = getattr(self, f'linear_{self._rank}')
        out = linear(x)
        # print(f"[linear_{self._rank}]input sum {x.sum()}, weight sum {linear.weight.sum()}, bias sum {linear.bias.sum()}, output sum {out.sum()}")
        return out
    
    def linear_weight(self):
        linear: nn.Linear = getattr(self, f'linear_{self._rank}')
        return linear.weight, linear.bias

def init_rpc(rank, host, port):
    os.environ['MASTER_ADDR'] = str(host)
    os.environ['MASTER_PORT'] = str(port)
    
    worker_name = f"work{rank}"
    
    # set options and set device map
    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=128)
    if rank == 0:
        # map rank_0 cuda:0 to other ranks
        for i in range(world_size):
            options.set_device_map(f"work{i}", {0 : i})

    rpc.init_rpc(worker_name, rank=rank, world_size=world_size, rpc_backend_options=options)

def run_worker(rank, host, port, backend):
    init_rpc(rank, host, port)
    
    device = 'cuda'
    if rank == 0:
        x = torch.randn([1000, 1000], device=device)

        # init rref of model of each tsage
        rref_rank0_model: PyRRef = rpc.remote("work0", TestModel, args=(0, 'cpu' if device == 'cpu' else 'cuda:0'))
        rref_rank1_model: PyRRef = rpc.remote("work1", TestModel, args=(1, 'cpu' if device == 'cpu' else 'cuda:1'))
        rref_rank2_model: PyRRef = rpc.remote("work2", TestModel, args=(2, 'cpu' if device == 'cpu' else 'cuda:2'))
        rref_rank3_model: PyRRef = rpc.remote("work3", TestModel, args=(3, 'cpu' if device == 'cpu' else 'cuda:3'))
        
        # init blocking proxy of each rref
        sync_proxy_model_0: RRefProxy = rref_rank0_model.rpc_sync()
        sync_proxy_model_1: RRefProxy = rref_rank1_model.rpc_sync()
        sync_proxy_model_2: RRefProxy = rref_rank2_model.rpc_sync()
        sync_proxy_model_3: RRefProxy = rref_rank3_model.rpc_sync()

        # init non blocking proxy of 
        async_proxy_model_0: RRefProxy = rref_rank0_model.rpc_async()
        async_proxy_model_1: RRefProxy = rref_rank1_model.rpc_async()
        async_proxy_model_2: RRefProxy = rref_rank2_model.rpc_async()
        async_proxy_model_3: RRefProxy = rref_rank3_model.rpc_async()
        
        # run async call 10 times
        for _ in range(10):
            out_0 = async_proxy_model_0.forward(x).wait()
            out_1 = async_proxy_model_1.forward(out_0).wait()
            out_2 = async_proxy_model_2.forward(out_1).wait()
            out_3 = async_proxy_model_3.forward(out_2).wait()
            print(f"async check sum {out_3.sum()}")
            
        # run sync call 10 times
        for _ in range(10):
            out_0: torch.Tensor = sync_proxy_model_0.forward(x)
            out_1: torch.Tensor = sync_proxy_model_1.forward(out_0)        
            out_2: torch.Tensor = sync_proxy_model_2.forward(out_1)
            out_3: torch.Tensor = sync_proxy_model_3.forward(out_2)
            print(f"sync check sum {out_3.sum()}")
        
    # barrier for rpc
    rpc.shutdown()

if __name__ == "__main__":
    host = 'localhost'
    port = '9001'
    backend = 'nccl'
    mp.spawn(run_worker, nprocs=world_size, args=(host, port, backend))
```


If `TestModel` is replaced with a smaller one, outputs are fine:

```python
class TestModel(nn.Module):
    def __init__(self, rank, device) -> None:
        super().__init__()
        self._rank = rank
        if rank == 0:
            self.linear_0 = nn.Linear(1000, 128).to(device)
        elif rank == 1:
            self.linear_1 = nn.Linear(128, 256).to(device)
        elif rank == 2:
            self.linear_2 = nn.Linear(256, 16).to(device)
        elif rank == 3:
            self.linear_3 = nn.Linear(16, 1).to(device)
```

outputs:
```
async check sum -195.64047241210938
async check sum -195.64047241210938
async check sum -195.64047241210938
async check sum -195.64047241210938
async check sum -195.64047241210938
async check sum -195.64047241210938
async check sum -195.64047241210938
async check sum -195.64047241210938
async check sum -195.64047241210938
async check sum -195.64047241210938
sync check sum -195.64047241210938
sync check sum -195.64047241210938
sync check sum -195.64047241210938
sync check sum -195.64047241210938
sync check sum -195.64047241210938
sync check sum -195.64047241210938
sync check sum -195.64047241210938
sync check sum -195.64047241210938
sync check sum -195.64047241210938
sync check sum -195.64047241210938
```

So, it seems that the bug may occur when the module is relevantly large.

---

### Versions

PyTorch version: 1.12.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun 22 2022, 20:18:18)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-43-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration:
GPU 0: NVIDIA A100 80GB PCIe
GPU 1: NVIDIA A100 80GB PCIe
GPU 2: NVIDIA A100 80GB PCIe
GPU 3: NVIDIA A100 80GB PCIe
GPU 4: NVIDIA A100 80GB PCIe
GPU 5: NVIDIA A100 80GB PCIe
GPU 6: NVIDIA A100 80GB PCIe
GPU 7: NVIDIA A100 80GB PCIe

Nvidia driver version: 515.65.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] colossalai==0.1.8+torch1.12cu11.3
[pip3] numpy==1.23.1
[pip3] torch==1.12.0+cu113
[pip3] torchaudio==0.12.0+cu113
[pip3] torchvision==0.13.0
[conda] Could not collect

cc @ezyang @gchanan @zou3519 @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501 @jjlilley @mrzzd