[TPU] Support collective communications in XLA devices #6813
Conversation
WoosukKwon
added
tpu
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge). To run full CI, you can do one of these:
🚀 |
|
This PR was part of #5871, but separated out to get a quick review. |
vllm/distributed/parallel_state.py
Outdated
| pynccl_comm: Optional[Any] # PyNccl communicator | ||
| ca_comm: Optional[Any] # Custom allreduce communicator | ||
| mq_broadcaster: Optional[Any] # shared memory broadcaster | ||
| use_xla: bool # Whether to use PyTorch XLA communicator |
There was a problem hiding this comment.
does tpu platform support NCCL? if not, creating these communicators might lead to error.
There was a problem hiding this comment.
TPU doesn't support NCCL, but I didn't see any error with the other communicators.
There was a problem hiding this comment.
The TPU backend uses gloo backend in addition to the distributed runtime in xm. Maybe that's the reason.
vllm/distributed/parallel_state.py
Outdated
| @@ -125,6 +129,7 @@ class GroupCoordinator: | |||
| pynccl_comm: Optional[Any] # PyNccl communicator | |||
| ca_comm: Optional[Any] # Custom allreduce communicator | |||
| mq_broadcaster: Optional[Any] # shared memory broadcaster | |||
| use_xla: bool # Whether to use PyTorch XLA communicator | |||
There was a problem hiding this comment.
use_xxx is a initialization parameter, and we usually hold communicator inside group coordinator.
Can you add a tpu_communicator under https://github.com/vllm-project/vllm/tree/main/vllm/distributed/device_communicators ?
One additional benefit, is that you can implement the gather logic to allgather, without intrusive change to logits_processor.py .
There was a problem hiding this comment.
One additional benefit, is that you can implement the gather logic to allgather, without intrusive change to logits_processor.py .
This is actually not the case because the TPU backend explicitly requires all-gather, which means each device's output should not be None. If we implement gather by using all-gather and outputting None for non-root ranks, XLA will raise an error.
| return | ||
| self.disabled = False | ||
|
|
||
| pjrt.initialize_multiprocess(local_rank, world_size) |
There was a problem hiding this comment.
you can get rank and world size from the group
There was a problem hiding this comment.
Thanks for letting me know! Updated the PR. PTAL.
There was a problem hiding this comment.
do you need to remove the code inside tpu worker? I don't know if pjrt and xm support initialization for multiple times.
There was a problem hiding this comment.
@youkaichao Good point. That's actually updated in #5871. In the current main branch, there's no code initializing the XLA's distributed runtime.
|
@youkaichao Thanks for your review and suggestions to the PR! |
…6813) Signed-off-by: Alvant <[email protected]>
This PR adds support for collective communications in XLA devices (TPU). It is simply implemented by falling back to
xm.all_reduceandxm.all_gatherfor TPU devices. One difference is the gather operation in logits processor, wheregatheris replaced byall-gatherto meet the SPMD restriction in XLA.