frozenleaves/vllm
1
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2025-10-20 14:53:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix weigit loading for GQA with TP (#2379)

Browse Source
This commit is contained in:
Chenhui Zhang
2024-01-16 07:43:59 +08:00
committed by GitHub
parent bfc072addf
commit f780504d12
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/model_executor/layers/linear.py
View File

@ -423,7 +423,10 @@ class QKVParallelLinear(ColumnParallelLinear):
shard_offset = shard_offset // param.pack_factor
param_data = param_data.narrow(output_dim, shard_offset,
shard_size)
shard_id = tp_rank // self.num_kv_head_replicas
if loaded_shard_id == "q":
shard_id = tp_rank
else:
shard_id = tp_rank // self.num_kv_head_replicas
start_idx = shard_id * shard_size
loaded_weight = loaded_weight.narrow(output_dim, start_idx,
shard_size)