[Request] Optimize gloo with oneCCL for CPUs

[jianan-gu]

Summary

To speed up LLM model inference on multiple Xeon CPUs, we require oneCCL to optimize the existing PyTorch native communication backend for Xeon CPUs (gloo).

Use case

Tensor Parallel is key for LLM models to run inference on multiple CPUs, and its performance relies on the
efficiency of communication backend implementations for communication ops like allreduce, allgather, gather and etc.

PyTorch communication backends:

• https://docs.pytorch.org/docs/stable/distributed.html

PyTorch + Transformers Tensor Parallel use case and status on CPUs:

• The performance of model with TP worse than without TP in CPU pytorch/pytorch#147596

As links listed above, the use case shows the current PyTorch native communication backend for CPU (gloo) is far from performant and needs to be optimized.

Proposed solution

Use the optimized communication ops from oneCCL like allreduce, allgather, gather and etc to replace the existing implementations in PyTorch gloo

Justification

This feature will help the OOB performance of LLM inference on Xeon CPUs using the PyTorch + Transformers ecosystem.

[iBaer]

@jianan-gu
Thanks for the interesting suggestion.
I have been reading through the attached Pytorch thread as well.
It would be good to understand a couple of things beforehand.

1. Do I understand your wording correctly? Would you like to see changes in gloo that are in sync with oneCCL collectives?
   1.1 If yes, does this mean no additional changes to oneCCL's code will be required?
2. What are the intended scenarios: scale up / scale out? How many ranks? Particular message sizes of large interest?
3. Similar to the pytorch thread, could we have a detailed overview of a benchmark that can be used for this discussion?