Repository: pytorch
Issue: [Distributed] test_dynamic_rpc_existing_rank_can_communicate_with_new_rank_cuda fails in caching allocator
Labels: oncall: distributed

### 🐛 Describe the bug

Code to reproduce the issue:
```python
python distributed/rpc/test_tensorpipe_agent.py -v -k test_dynamic_rpc_existing_rank_can_communicate_with_new_rank_cuda
```
Output:
```python
INFO:torch.distributed.nn.jit.instantiator:Created a temporary directory at /tmp/tmpce6ewbo1
INFO:torch.distributed.nn.jit.instantiator:Writing /tmp/tmpce6ewbo1/_remote_module_non_scriptable.py
test_dynamic_rpc_existing_rank_can_communicate_with_new_rank_cuda (__main__.TensorPipeTensorPipeAgentRpcTest) ... INFO:numba.cuda.cudadrv.driver:init
INFO:torch.testing._internal.common_distributed:Started process 0 with pid 64064
INFO:torch.testing._internal.common_distributed:Started process 1 with pid 64065
INFO:torch.testing._internal.common_distributed:Started process 2 with pid 64066
INFO:torch.testing._internal.common_distributed:Started process 3 with pid 64067
INFO:torch.distributed.nn.jit.instantiator:Created a temporary directory at /tmp/tmpj5jtkv1l
INFO:torch.distributed.nn.jit.instantiator:Writing /tmp/tmpj5jtkv1l/_remote_module_non_scriptable.py
INFO:torch.distributed.nn.jit.instantiator:Created a temporary directory at /tmp/tmp3xq_65qx
INFO:torch.distributed.nn.jit.instantiator:Writing /tmp/tmp3xq_65qx/_remote_module_non_scriptable.py
INFO:torch.distributed.nn.jit.instantiator:Created a temporary directory at /tmp/tmp0oc80wck
INFO:torch.distributed.nn.jit.instantiator:Writing /tmp/tmp0oc80wck/_remote_module_non_scriptable.py
INFO:torch.distributed.nn.jit.instantiator:Created a temporary directory at /tmp/tmpp47jqr8l
INFO:torch.distributed.nn.jit.instantiator:Writing /tmp/tmpp47jqr8l/_remote_module_non_scriptable.py
INFO:torch.testing._internal.common_distributed:Starting event listener thread for rank 1
INFO:torch.testing._internal.common_distributed:Starting event listener thread for rank 0
INFO:torch.testing._internal.common_distributed:Starting event listener thread for rank 2
INFO:torch.testing._internal.common_distributed:Starting event listener thread for rank 3
INFO:torch.distributed.distributed_c10d:Added key: store_based_barrier_key:1 to store for rank: 0
INFO:torch.distributed.distributed_c10d:Added key: store_based_barrier_key:1 to store for rank: 1
INFO:torch.distributed.distributed_c10d:Added key: store_based_barrier_key:1 to store for rank: 3
INFO:torch.distributed.distributed_c10d:Rank 3: Completed store-based barrier for key:store_based_barrier_key:1 with 4 nodes.
INFO:torch.distributed.distributed_c10d:Added key: store_based_barrier_key:1 to store for rank: 2
INFO:torch.distributed.distributed_c10d:Rank 1: Completed store-based barrier for key:store_based_barrier_key:1 with 4 nodes.
INFO:torch.distributed.distributed_c10d:Rank 2: Completed store-based barrier for key:store_based_barrier_key:1 with 4 nodes.
INFO:torch.distributed.distributed_c10d:Rank 0: Completed store-based barrier for key:store_based_barrier_key:1 with 4 nodes.
terminate called after throwing an instance of 'c10::Error'
  what():  0 <= device && static_cast<size_t>(device) < device_allocator.size() INTERNAL ASSERT FAILED at "/workspace/src/pytorch/c10/cuda/CUDACachingAllocator.cpp":1602, please report a bug to PyTorch. Allocator not initialized for device 1: did you call init?
Exception raised from malloc at /workspace/src/pytorch/c10/cuda/CUDACachingAllocator.cpp:1602 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0x6c (0x7f63d94b96cc in /workspace/src/pytorch/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0xfa (0x7f63d948f12c in /workspace/src/pytorch/torch/lib/libc10.so)
frame #2: c10::detail::torchInternalAssertFail(char const*, char const*, unsigned int, char const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0x53 (0x7f63d94b73e3 in /workspace/src/pytorch/torch/lib/libc10.so)
...
```

Unsure, if the test is executed in an unsupported way or why CI isn't seeing this issue.

### Versions

Current master build at [fe7262329c](https://github.com/pytorch/pytorch/commit/fe7262329c6ae702df185bdbeac702e0d2edc123) on a node with 8x A100.



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