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 8, 20:15 UTC
[Minor] Sort safetensors files to ensure deterministic loading order (#33491)
Signed-off-by: Lihao Ran <imlihao.ran@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Lumosis committed
Feb 5, 2026 at 14:05 UTC
42d5d705f93b254179e062003e8504fbe04f1b30
1 file changed
+11
-2
vllm/model_executor/model_loader/weight_utils.py
+11
-2
@@ -19,6 +19,7 @@ from typing import IO, Any
19
import filelock
20
import huggingface_hub.constants
21
import numpy as np
22
+import regex as re
23
import torch
24
from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
25
from safetensors.torch import load, load_file, safe_open, save_file
@@ -143,6 +144,15 @@ def atomic_writer(
144
os.remove(temp_path)
145
146
147
+def _natural_sort_key(filepath: str) -> list:
148
+ """Natural sort key for filenames with numeric components, such as
149
+ model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors']"""
150
+ return [
151
+ int(s) if s.isdigit() else s
152
+ for s in re.split(r"(\d+)", os.path.basename(filepath))
153
+ ]
154
+
155
+
156
def maybe_download_from_modelscope(
157
model: str,
158
revision: str | None = None,
@@ -682,9 +692,8 @@ def safetensors_weights_iterator(
692
loading_desc += " (eager)"
693
694
leftover_state_dict: dict[str, torch.Tensor] = {}
685
-
695
for st_file in tqdm(
687
- hf_weights_files,
696
+ sorted(hf_weights_files, key=_natural_sort_key),
697
desc=loading_desc,
698
disable=not enable_tqdm(use_tqdm_on_load),
699
bar_format=_BAR_FORMAT,