A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 05:00 UTC
Filter safetensors files to download if .safetensors.index.json exists (#30537)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Dec 18, 2025 at 09:51 UTC
100f93d2bea44097916c568a7ac642e7fee915ed
1 file changed
+29
-6
vllm/model_executor/model_loader/weight_utils.py
+29
-6
@@ -23,6 +23,7 @@ import torch
23
from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
24
from safetensors.torch import load, load_file, safe_open, save_file
25
from tqdm.auto import tqdm
26
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
27
28
from vllm import envs
29
from vllm.config import ModelConfig
@@ -448,12 +449,31 @@ def download_weights_from_hf(
449
fs = HfFileSystem()
450
file_list = fs.ls(model_name_or_path, detail=False, revision=revision)
451
451
- # Use the first pattern found in the HF repo's files.
452
- for pattern in allow_patterns:
453
- matching = fnmatch.filter(file_list, pattern)
454
- if len(matching) > 0:
455
- allow_patterns = [pattern]
456
- break
452
+ # If downloading safetensors and an index file exists, use the
453
+ # specific file names from the index to avoid downloading
454
+ # unnecessary files (e.g., from subdirectories like "original/").
455
+ index_file = f"{model_name_or_path}/{SAFE_WEIGHTS_INDEX_NAME}"
456
+ if "*.safetensors" in allow_patterns and index_file in file_list:
457
+ index_path = hf_hub_download(
458
+ repo_id=model_name_or_path,
459
+ filename=SAFE_WEIGHTS_INDEX_NAME,
460
+ cache_dir=cache_dir,
461
+ revision=revision,
462
+ )
463
+ with open(index_path) as f:
464
+ weight_map = json.load(f)["weight_map"]
465
+ if weight_map:
466
+ # Extra [] so that weight_map files are treated as a
467
+ # single allow_pattern in the loop below
468
+ allow_patterns = [list(set(weight_map.values()))] # type: ignore[list-item]
469
+ else:
470
+ allow_patterns = ["*.safetensors"]
471
+ else:
472
+ # Use the first pattern found in the HF repo's files.
473
+ for pattern in allow_patterns:
474
+ if fnmatch.filter(file_list, pattern):
475
+ allow_patterns = [pattern]
476
+ break
477
except Exception as e:
478
logger.warning(
479
"Failed to get file list for '%s'. Trying each pattern in "
@@ -480,6 +500,9 @@ def download_weights_from_hf(
500
)
501
# If we have downloaded weights for this allow_pattern,
502
# we don't need to check the rest.
503
+ # allow_pattern can be a list (from weight_map) or str (glob)
504
+ if isinstance(allow_pattern, list):
505
+ break
506
if any(Path(hf_folder).glob(allow_pattern)):
507
break
508
time_taken = time.perf_counter() - start_time