[Distributed][refactor] Add base class for device-specific communicator #11324
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can do one of these:
🚀 |
MengqingCao
force-pushed
the
communicator
branch
from
c4f0481 to
eeb5aae
Compare
MengqingCao
force-pushed
the
communicator
branch
from
eeb5aae to
6e6501a
Compare
MengqingCao
force-pushed
the
communicator
branch
3 times, most recently
from
1a977d3 to
6de2b98
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
| class CommunicatorABC(ABC): | ||
| """ | ||
| CommunicatorBase ABC | ||
| """ | ||
|
|
||
| @abstractmethod | ||
| def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: | ||
| raise NotImplementedError | ||
|
|
||
| def gather(self, | ||
| input_: torch.Tensor, | ||
| dst: int = 0, | ||
| dim: int = -1) -> Optional[torch.Tensor]: | ||
|
|
||
| raise NotImplementedError | ||
|
|
||
| def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: | ||
| raise NotImplementedError | ||
|
|
||
|
|
||
| class CommunicatorBase(CommunicatorABC): |
There was a problem hiding this comment.
why are these two classes instead of one?
There was a problem hiding this comment.
The abstract class is just to provide a unified interface and indicate the object type, like following
vllm/vllm/distributed/parallel_state.py
Line 229 in 6de2b98
There was a problem hiding this comment.
@youkaichao WDYT of we keep both CommunicatorABC and CommunicatorBase here? Plz let me know if you have any suggestion on this.
There was a problem hiding this comment.
update: now we using CommunicatorProtocol instead of CommunicatorABC for type indication
MengqingCao
force-pushed
the
communicator
branch
from
b5d2063 to
f03eedb
Compare
MengqingCao
force-pushed
the
communicator
branch
from
242fb40 to
b085f82
Compare
MengqingCao
force-pushed
the
communicator
branch
5 times, most recently
from
cadfb32 to
cc6d46a
Compare
MengqingCao
force-pushed
the
communicator
branch
from
cc6d46a to
464594a
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
MengqingCao
force-pushed
the
communicator
branch
from
464594a to
1e986a0
Compare
MengqingCao
force-pushed
the
communicator
branch
from
1e986a0 to
79a5eb0
Compare
Signed-off-by: Mengqing Cao <[email protected]>
MengqingCao
force-pushed
the
communicator
branch
from
79a5eb0 to
6851cd0
Compare
|
CI failed due to network issues. This pr is ready for review now, thanks in advance! @youkaichao |
There was a problem hiding this comment.
@youkaichao Would you mind taking another look?
Or if you are worried about the code changes being too big and want us to split the PR, for example:
- a separate PR for CommunicatorBase and interface change
- Adapts cuda/rocm, hpu, tpu, xpu separately and split to 3 followup PRs
Please let us know, we'd like to do so.
|
sorry I'm super busy recently. will review this week. |
part of #11162
This PR provide a base class
CommunicatorBasefor device-specific communicators (HpuCommunicator,TpuCommunicatorandXpuCommunicator), avoiding the cumbersome dispatch in each communicator operator ofGroupCoordinator, e.g.,https://github.com/vllm-project/vllm/blob/main/vllm/distributed/parallel_state.py#L342-L353
In this pr, the communication-related classes are organized as the following fig. This allows new backends to implement their own communicators and dynamic dispatch them in the platform.
