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

Filter safetensors files to download if .safetensors.index.json exists (#30537)

Browse Source 
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Michael Goin 2025-12-18 09:51:17 -05:00 committed by GitHub
parent 96bf50a2c0
commit 100f93d2be
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 29 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
vllm/model_executor/model_loader/weight_utils.py
View File

@ -23,6 +23,7 @@ import torch
from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
from safetensors.torch import load, load_file, safe_open, save_file from safetensors.torch import load, load_file, safe_open, save_file
from tqdm.auto import tqdm from tqdm.auto import tqdm
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
from vllm import envs from vllm import envs
from vllm.config import ModelConfig from vllm.config import ModelConfig
@ -448,12 +449,31 @@ def download_weights_from_hf(
fs = HfFileSystem() fs = HfFileSystem()
file_list = fs.ls(model_name_or_path, detail=False, revision=revision) file_list = fs.ls(model_name_or_path, detail=False, revision=revision)
# Use the first pattern found in the HF repo's files. # If downloading safetensors and an index file exists, use the
for pattern in allow_patterns: # specific file names from the index to avoid downloading
matching = fnmatch.filter(file_list, pattern) # unnecessary files (e.g., from subdirectories like "original/").
if len(matching) > 0: index_file = f"{model_name_or_path}/{SAFE_WEIGHTS_INDEX_NAME}"
allow_patterns = [pattern] if "*.safetensors" in allow_patterns and index_file in file_list:
break index_path = hf_hub_download(
repo_id=model_name_or_path,
filename=SAFE_WEIGHTS_INDEX_NAME,
cache_dir=cache_dir,
revision=revision,
)
with open(index_path) as f:
weight_map = json.load(f)["weight_map"]
if weight_map:
# Extra [] so that weight_map files are treated as a
# single allow_pattern in the loop below
allow_patterns = [list(set(weight_map.values()))] # type: ignore[list-item]
else:
allow_patterns = ["*.safetensors"]
else:
# Use the first pattern found in the HF repo's files.
for pattern in allow_patterns:
if fnmatch.filter(file_list, pattern):
allow_patterns = [pattern]
break
except Exception as e: except Exception as e:
logger.warning( logger.warning(
"Failed to get file list for '%s'. Trying each pattern in " "Failed to get file list for '%s'. Trying each pattern in "
@ -480,6 +500,9 @@ def download_weights_from_hf(
) )
# If we have downloaded weights for this allow_pattern, # If we have downloaded weights for this allow_pattern,
# we don't need to check the rest. # we don't need to check the rest.
# allow_pattern can be a list (from weight_map) or str (glob)
if isinstance(allow_pattern, list):
break
if any(Path(hf_folder).glob(allow_pattern)): if any(Path(hf_folder).glob(allow_pattern)):
break break
time_taken = time.perf_counter() - start_time time_taken = time.perf_counter() - start_time