Repository: pytorch
Issue: [RFC] UCC integration in ProcessGroupNCCL
Labels: oncall: distributed,feature,triaged

This RFC proposes to use https://github.com/facebookresearch/torch_ucc as a fallback path for `ProcessGroupNCCL` to support cases where NCCL doesn't support.
For now, these cases will be just communications on CPU tensors. Feel free to leave a comment on this issue or PRs if you have any opinion.

## How does the integration look like?

As demonstrated in https://github.com/pytorch/pytorch/pull/69552, `torch_ucc.so` will NOT be linked to PyTorch at build time.
Instead, `ProcessGroupNCCL` will load the `torch_ucc.so` at runtime when the user creates an instance of `ProcessGroupNCCL` for the first time.
The handle of the loaded library will be stored as a static member `uccLib_` of `ProcessGroupNCCL`.
The `uccLib_` will be the interface for communications between PyTorch and `torch_ucc.so`.
Dynamic loading of `torch_ucc.so` will only happen if the user set an environment variable `TORCH_UCC_LIBRARY_PATH` specifying where the `torch_ucc.so` is located.
There will be a cmake flag `USE_NCCL_WITH_UCC` controlling whether to enable this UCC integration at build time.
`USE_NCCL_WITH_UCC` will default to `ON` because there won't be any user-visible changes unless the user manually specifies `TORCH_UCC_LIBRARY_PATH`.
In the future, we might make PyTorch smarter to automatically find `torch_ucc.so` installed in the user's system, but I believe a better time to do so is after all PyTorch unit tests (discussed below) pass with UCC.
The `torch_ucc.so` will be loaded with `RTLD_LOCAL` so that symbols in `torch_ucc.so` will not leak to PyTorch.

As demonstrated in https://github.com/pytorch/pytorch/pull/69564, whenever `uccLib_` is valid, creating an instance of `ProcessGroupNCCL` will always lead to a creation of `ProcessGroupUCC`.
Pointer to `ProcessGroupUCC` instance will be stored as a member of `ProcessGroupNCCL` named `uccPG_`.

As demonstrated in https://github.com/pytorch/pytorch/pull/69587, the dispatch to UCC can be implemented like this
```C++
c10::intrusive_ptr<ProcessGroup::Work> ProcessGroupNCCL::allgather(
    std::vector<std::vector<at::Tensor>>& outputTensors,
    std::vector<at::Tensor>& inputTensors,
    const AllgatherOptions& opts) {
  if (outputTensors[0].device().type() != at::DeviceType::CUDA && uccPG_ != nullptr) {
    return uccPG_->allgather(outputTensors, inputTensors, opts);
  }
  // original implementation are kept unchanged.
}
```
Enablement of the UCC path should be done incrementally. Ideally, there will be one PR for each operation. PRs enabling the UCC path for an operation should add the dispatch code and enable unit tests through the NCCL PG in the same PR.

## How will this integration be tested?

The integration will be tested mainly by PyTorch's unit tests.
Today, PyTorch's unit tests for `torch.distributed` contains CPU tests running on MPI and Gloo, and GPU tests running on NCCL and Gloo.
We want to make sure these CPU tests run successfully through the NCCL PG. At the beginning, all tests will be skipped through the NCCL PG, but will be enabled incrementally. Our goal is to eventually enable all of them in CI.

Both PyTorch CI and torch_ucc CI will run these tests. https://github.com/pytorch/pytorch/pull/71017 prepares PyTorch CI for these tests, and https://github.com/facebookresearch/torch_ucc/pull/58 prepares `torch_ucc` CI for the tests. These two PRs are WIP, but reviews are welcome.

## Other questions

There are many other important questions that are not covered here in this RFC. I believe a better time to answer these questions is after we successfully made all PyTorch unit tests running on UCC. These questions include but are not limited to:
- What standard should we achieve in order to officially announce the UCC integration to users? What else should we test?
- Instead of manually setting `TORCH_UCC_LIBRARY_PATH`, how can we make the build easier for users?
- When should we evaluate packaging `libucc.so`, `libucs.so`, `libucp.so`, etc. into the PyTorch wheels?

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang