From 15713e3b7579d56758fab1150c99dd49633b5669 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 21 Oct 2024 22:14:29 +0100 Subject: [PATCH] [BugFix] Update draft model TP size check to allow matching target TP size (#9394) Co-authored-by: Baoyuan Qi --- vllm/config.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/config.py b/vllm/config.py index f57aa404..00dd047e 100644 --- a/vllm/config.py +++ b/vllm/config.py @@ -1408,11 +1408,11 @@ class SpeculativeConfig: else: speculative_draft_tensor_parallel_size = \ target_parallel_config.tensor_parallel_size - elif speculative_draft_tensor_parallel_size != 1: - # TODO(wooyeon): allow tp values larger than 1 + elif speculative_draft_tensor_parallel_size not in ( + 1, target_parallel_config.tensor_parallel_size): raise ValueError( f"{speculative_draft_tensor_parallel_size=} cannot be " - f"other value than 1") + f"other value than 1 or target model tensor_parallel_size") draft_parallel_config = ParallelConfig( pipeline_parallel_size=target_parallel_config.