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

Signed-off-by: wenxindongwork <wenxindong@google.com>
This commit is contained in:
wenxindongwork 2025-11-01 11:14:44 -06:00 committed by GitHub
parent 99d69af9ec
commit af6e19f50f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -67,6 +67,7 @@ from vllm.outputs import (
RequestOutput,
ScoringRequestOutput,
)
from vllm.platforms import current_platform
from vllm.pooling_params import PoolingParams
from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
from vllm.tasks import PoolingTask
@ -289,7 +290,11 @@ class LLM:
# warn about single-process data parallel usage.
_dp_size = int(kwargs.get("data_parallel_size", 1))
_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(
f"LLM(data_parallel_size={_dp_size}) is not supported for single-"
"process usage and may hang. Please use "