Skip to content

Commit 5bce45e

Browse files
authored
[TPU] prevent single-process DP (#1593)
Signed-off-by: AlpinDale <alpindale@gmail.com>
1 parent 2c3b58c commit 5bce45e

1 file changed

Lines changed: 2 additions & 1 deletion

File tree

aphrodite/endpoints/llm.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
RequestOutput,
4949
ScoringRequestOutput,
5050
)
51+
from aphrodite.platforms import current_platform
5152
from aphrodite.quantization import QuantizationMethods
5253
from aphrodite.tasks import PoolingTask
5354
from aphrodite.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer, get_cached_tokenizer
@@ -252,7 +253,7 @@ def __init__(
252253
# warn about single-process data parallel usage.
253254
_dp_size = int(kwargs.get("data_parallel_size", 1))
254255
_distributed_executor_backend = kwargs.get("distributed_executor_backend")
255-
if _dp_size > 1 and not _distributed_executor_backend == "external_launcher":
256+
if _dp_size > 1 and not _distributed_executor_backend == "external_launcher" and not current_platform.is_tpu():
256257
raise ValueError(
257258
f"LLM(data_parallel_size={_dp_size}) is not supported for single-"
258259
"process usage and may hang. Please use "

0 commit comments

Comments
 (0)