xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-13 20:55:37 +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
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

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