xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-17 07:35:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] Fix num_heads value for simple connector when tp enabled (#12074)

Browse Source 
Signed-off-by: Shangming Cai <caishangming@linux.alibaba.com>
This commit is contained in:
shangmingc 2025-01-20 10:56:43 +08:00 committed by GitHub
parent bbe5f9de7d
commit df450aa567
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
vllm/distributed/kv_transfer/kv_connector/simple_connector.py
View File

@ -35,6 +35,7 @@ class SimpleConnector(KVConnectorBase):
): ):
self.config = config.kv_transfer_config self.config = config.kv_transfer_config
self.tp_size = config.parallel_config.tensor_parallel_size
if self.config.kv_connector == "PyNcclConnector": if self.config.kv_connector == "PyNcclConnector":
from vllm.distributed.kv_transfer.kv_pipe.pynccl_pipe import ( from vllm.distributed.kv_transfer.kv_pipe.pynccl_pipe import (
@ -161,7 +162,7 @@ class SimpleConnector(KVConnectorBase):
end_layer = model_executable.model.end_layer end_layer = model_executable.model.end_layer
model_config = model_executable.model.config model_config = model_executable.model.config
num_heads = model_config.num_key_value_heads num_heads = int(model_config.num_key_value_heads / self.tp_size)
hidden_size = model_config.hidden_size hidden_size = model_config.hidden_size
num_attention_heads = model_config.num_attention_heads num_attention_heads = model_config.num_attention_heads
head_size = int(hidden_size / num_attention_heads) head_size = int(hidden_size / num_attention_heads)