A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[TPU] optimize the all-reduce performance (#15903)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Apr 2, 2025 at 17:25 UTC
01b6113659cf4ad115c39e2ebb18de7535e423a7
3 files changed
+16
-2
vllm/distributed/device_communicators/tpu_communicator.py
+6
-1
@@ -22,6 +22,8 @@ if current_platform.is_tpu():
22
import torch_xla.core.xla_model as xm
23
import torch_xla.runtime as xr
24
from torch_xla._internal import pjrt
25
+ from torch_xla.distributed.xla_multiprocessing import (
26
+ create_optimized_replica_groups)
27
28
if USE_RAY:
29
from vllm.executor import ray_utils
@@ -79,9 +81,12 @@ class TpuCommunicator(DeviceCommunicatorBase):
81
82
pjrt.initialize_multiprocess(local_rank, local_world_size)
83
xr._init_world_size_ordinal()
84
+ self.groups = create_optimized_replica_groups()
85
86
def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
84
- return xm.all_reduce(xm.REDUCE_SUM, input_)
87
+ # TODO: Remove the groups specification after XLA compiler can support
88
+ # auto-reordering the ring order for all-reduce.
89
+ return xm.all_reduce(xm.REDUCE_SUM, input_, groups=self.groups)
90
91
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
92
assert dim == -1, "TPUs only support dim=-1 for all-gather."
vllm/distributed/parallel_state.py
+4
-1
@@ -119,11 +119,13 @@ def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
119
120
121
if supports_custom_op():
122
+ from vllm.platforms import current_platform
123
direct_register_custom_op(
124
op_name="all_reduce",
125
op_func=all_reduce,
126
mutates_args=[],
127
fake_impl=all_reduce_fake,
128
+ dispatch_key=current_platform.dispatch_key,
129
)
130
131
@@ -219,7 +221,8 @@ class GroupCoordinator:
221
self.cpu_group, 1 << 22, 6)
222
223
from vllm.platforms import current_platform
222
- self.use_custom_op_call = current_platform.is_cuda_alike()
224
+ self.use_custom_op_call = (current_platform.is_cuda_alike()
225
+ or current_platform.is_tpu())
226
227
@property
228
def first_rank(self):
vllm/v1/worker/tpu_worker.py
+6
@@ -84,6 +84,12 @@ class TPUWorker:
84
85
def init_device(self):
86
os.environ["PJRT_DEVICE"] = "TPU"
87
+ # Note: Currently the XLA compiler wrongly uses 2D ring strategy on 1D
88
+ # ring, the xla tpu compiler flag
89
+ # `xla_tpu_force_1d_allreduce_at_chunk_count` is a temporary solution to
90
+ # fix this. It will be removed after the bug in XLA compiler is fixed.
91
+ os.environ["LIBTPU_INIT_ARGS"] = (
92
+ "--xla_tpu_force_1d_allreduce_at_chunk_count=1")
93
torch.set_grad_enabled(False)
94
torch.set_default_dtype(self.model_config.dtype)
95