xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 10:40:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

use get_tensor in safe_open (#1696)

Browse Source
This commit is contained in:
twaka 2023-11-19 09:45:18 +09:00 committed by GitHub
parent edb305584b
commit e946260cf3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
vllm/model_executor/weight_utils.py
View File

@ -243,8 +243,8 @@ def hf_model_weights_iterator(
for st_file in hf_weights_files:
with safe_open(st_file, framework="pt") as f:
for name in f.keys():
param = f.get_slice(name)
yield name, convert_pyslice_to_tensor(param)
param = f.get_tensor(name)
yield name, param
else:
for bin_file in hf_weights_files:
state = torch.load(bin_file, map_location="cpu")
@ -265,12 +265,7 @@ def convert_pyslice_to_tensor(x: Any) -> torch.Tensor:
tensor first.
"""
if not isinstance(x, torch.Tensor):
try:
x = x[:]
except IndexError:
# IndexError happens when the tensor is empty.
# transformer.h.0.attn.masked_bias is empty in some gpt2 models.
return torch.Tensor()
x = x[:]
return x