Unverified Commit af6e19f5 authored by wenxindongwork's avatar wenxindongwork Committed by GitHub
Browse files

[Core][TPU] Support TPU Data Parallalism (#27365)


Signed-off-by: default avatarwenxindongwork <wenxindong@google.com>
parent 99d69af9
...@@ -67,6 +67,7 @@ from vllm.outputs import ( ...@@ -67,6 +67,7 @@ from vllm.outputs import (
RequestOutput, RequestOutput,
ScoringRequestOutput, ScoringRequestOutput,
) )
from vllm.platforms import current_platform
from vllm.pooling_params import PoolingParams from vllm.pooling_params import PoolingParams
from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
from vllm.tasks import PoolingTask from vllm.tasks import PoolingTask
...@@ -289,7 +290,11 @@ class LLM: ...@@ -289,7 +290,11 @@ class LLM:
# warn about single-process data parallel usage. # warn about single-process data parallel usage.
_dp_size = int(kwargs.get("data_parallel_size", 1)) _dp_size = int(kwargs.get("data_parallel_size", 1))
_distributed_executor_backend = kwargs.get("distributed_executor_backend") _distributed_executor_backend = kwargs.get("distributed_executor_backend")
if _dp_size > 1 and not _distributed_executor_backend == "external_launcher": if (
_dp_size > 1
and not _distributed_executor_backend == "external_launcher"
and not current_platform.is_tpu()
):
raise ValueError( raise ValueError(
f"LLM(data_parallel_size={_dp_size}) is not supported for single-" f"LLM(data_parallel_size={_dp_size}) is not supported for single-"
"process usage and may hang. Please use " "process usage and may hang. Please use "
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment