Skip to content

GitLab

Unverified Commit 15713e3b authored by Nick Hill's avatar Nick Hill Committed by GitHub
Browse files

[BugFix] Update draft model TP size check to allow matching target TP size (#9394) 


Co-authored-by: default avatarBaoyuan Qi <qibaoyuan@126.com>
parent d621c43d
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
vllm/config.py
View file @ 15713e3b
...@@ -1408,11 +1408,11 @@ class SpeculativeConfig: ...@@ -1408,11 +1408,11 @@ class SpeculativeConfig:
else: else:
speculative_draft_tensor_parallel_size = \ speculative_draft_tensor_parallel_size = \
target_parallel_config.tensor_parallel_size target_parallel_config.tensor_parallel_size
elif speculative_draft_tensor_parallel_size != 1: elif speculative_draft_tensor_parallel_size not in (
# TODO(wooyeon): allow tp values larger than 1 1, target_parallel_config.tensor_parallel_size):
raise ValueError( raise ValueError(
f"{speculative_draft_tensor_parallel_size=} cannot be " 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( draft_parallel_config = ParallelConfig(
pipeline_parallel_size=target_parallel_config. pipeline_parallel_size=target_parallel_config.
......
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