Repository: pytorch
Issue: distributed: prefixstore doesn't recognize custom store
Labels: oncall: distributed

### 🐛 Describe the bug

I am using a custom python based store to address the issue of torch's `TCPStore` running out of file handle resources when running on a large number of GPUs. However, torch by default wraps the store into a `PrefixStore` which seems to disable the python API.

I am getting the following error
```
L ASSERT FAILED at "/gpfs/alpine2/bif143/proj-shared/summit_py310/conda-bld/pytorch-base_1708839867265/work/torch/csrc/distributed/c10d/init.cpp":151, please report a bug to PyTorch. Not implemented.
Exception raised from get at /gpfs/alpine2/bif143/proj-shared/summit_py310/conda-bld/pytorch-base_1708839867265/work/torch/csrc/distributed/c10d/init.cpp:151 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0xfc (0x200072d39d3c in /gpfs/alpine2/proj-shared/bif143/summit_py310/lib/python3.10/site-packages/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&) + 0x124 (0x200072cd1d04 in /gpfs/alpine2/proj-shared/bif143/summit_py310/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #2: c10::detail::torchInternalAssertFail(char const*, char const*, unsigned int, char const*, char const*) + 0x70 (0x200072d363e0 in /gpfs/alpine2/proj-shared/bif143/summit_py310/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #3: <unknown function> + 0xe74814 (0x2000429a4814 in /gpfs/alpine2/proj-shared/bif143/summit_py310/lib/python3.10/site-packages/torch/lib/libtorch_python.so)
frame #4: c10d::PrefixStore::get(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0x6c (0x20004993f4bc in /gpfs/alpine2/proj-shared/bif143/summit_py310/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #5: c10d::PrefixStore::get(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0x6c (0x20004993f4bc in /gpfs/alpine2/proj-shared/bif143/summit_py310/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #6: c10d::ProcessGroupNCCL::broadcastUniqueNCCLID(ncclUniqueId*, bool, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, int) + 0x234 (0x20004df5e7e4 in /gpfs/alpine2/proj-shared/bif143/summit_py310/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #7: c10d::ProcessGroupNCCL::getNCCLComm(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, std::vector<c10::Device, std::allocator<c10::Device> > const&, c10d::OpType, int, bool) + 0x294 (0x20004df64a44 in /gpfs/alpine2/proj-shared/bif143/summit_py310/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
```

implying that the torch can't resolve the virtual overload of the `Store->get()` function, even though it is implemented in the custom store.

My store:
```
from mpi4py import MPI
import torch.distributed as dist

class MPIStore(dist.Store):
    def __init__(self, comm=None):
        super().__init__()
        self.comm = comm
        if self.comm is None:
            self.comm = MPI.COMM_WORLD

    # Function to get value associated with a key
    def get(self, key):
        return self.comm.bcast(None)

    # Function to put a key-value pair into the store
    def set(self, key, value):
        self.comm.bcast(value)

    def add(self, key, value):
        return self.comm.Get_size()

    def wait(self, keys, timeout):
        pass

if __name__ == '__main__':
    dist._test_python_store(dist.PrefixStore('prefix', MPIStore()))
```

(i.e. just replace the complicated key/value store logic by a simple brodcast to set up NCCL)

Current output of this reproducer:
```
$ jsrun -n 1 python mpi_store.py 
Warning: OMP_NUM_THREADS=16 is greater than available PU's
Traceback (most recent call last):
  File "/gpfs/alpine2/bif143/proj-shared/glaser/affinity_clip/affinity_clip/mpi_store.py", line 26, in <module>
    dist._test_python_store(dist.PrefixStore('prefix', MPIStore()))
RuntimeError: Tried to call pure virtual function "::c10d::Store::add"
```

Expected output:
``` 
TypeError: Object of type 'NoneType' is not an instance of 'bytes'
``` 
(yes, that is an error, too, but it is expected from running this somewhat hacky class in a unit test setting rather than inside the distributed communication setup)

Workaround:
```
torch.distributed.distributed_c10d.PrefixStore = lambda x, s: s
torch.distributed.init_process_group("nccl",
                                     store=MPIStore(),
                                     rank=comm.Get_rank(),
                                     world_size=comm.Get_size())
```

i.e. monkey patch the PrefixStore to be a no-op.

It would be nice if PrefixStore could be updated to accept custom Python stores.

To address the underlying issue of distributed setup, it should be possible to initialize NCCL using MPI communicators instead of using TCP. Perhaps revamping the logic behind the k/v stores and reducing them to a `bcast` call might help here, if it applies to all other back ends as well.

### Versions

tested with 2.0

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225