From 8a25d3a71ac8b0b06c37935ed67e7c35aa901bf5 Mon Sep 17 00:00:00 2001 From: YingchaoX Date: Fri, 19 Jan 2024 01:39:46 +0800 Subject: [PATCH] fix stablelm.py tensor-parallel-size bug (#2482) --- vllm/model_executor/models/stablelm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/stablelm.py b/vllm/model_executor/models/stablelm.py index de5e246021b3b..cf842d087669f 100644 --- a/vllm/model_executor/models/stablelm.py +++ b/vllm/model_executor/models/stablelm.py @@ -99,7 +99,7 @@ class StablelmAttention(nn.Module): self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_key_value_heads * self.head_dim - if (self.head_dim * self.num_heads) != self.hidden_size: + if (self.head_dim * self.num_heads * tp_size) != self.hidden_size: raise ValueError( f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" f" and `num_heads`: {self.num_heads}).")