fix: correct load_weights track logic and enable weight integrity for… (#41811)
Signed-off-by: Yipeng Hu <i26268@metax-tech.com> Signed-off-by: HuYiPeng <144002351+MynameFelix@users.noreply.github.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Yipeng Hu <i26268@metax-tech.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
HuYiPeng committed
Jul 10, 2026 at 22:08 UTC
85c09e9885e346ea1612da30ebff5a75f67d2350
2 files changed
+87
tests/model_executor/model_loader/test_filter_duplicate_safetensors.py
new
+79
@@ -0,0 +1,79 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import json
5
+import os
6
+import tempfile
7
+
8
+import pytest
9
+
10
+from vllm.model_executor.model_loader.weight_utils import (
11
+ filter_duplicate_safetensors_files,
12
+)
13
+
14
+
15
+def test_filter_duplicate_safetensors_files_missing_weight():
16
+ with tempfile.TemporaryDirectory() as tmpdir:
17
+ existing_file = os.path.join(tmpdir, "model-00001-of-00002.safetensors")
18
+ with open(existing_file, "wb") as f:
19
+ f.write(b"")
20
+
21
+ existing_file2 = os.path.join(tmpdir, "model-00002-of-00002.safetensors")
22
+ with open(existing_file2, "wb") as f:
23
+ f.write(b"")
24
+
25
+ index_file = os.path.join(tmpdir, "model.safetensors.index.json")
26
+ index_content = {
27
+ "weight_map": {
28
+ "layer.0.weight": "model-00001-of-00002.safetensors",
29
+ "layer.1.weight": "model-00002-of-00002.safetensors",
30
+ "layer.2.weight": "model-00003-of-00002.safetensors",
31
+ }
32
+ }
33
+ with open(index_file, "w") as f:
34
+ json.dump(index_content, f)
35
+
36
+ hf_weights_files = [
37
+ os.path.join(tmpdir, "model-00001-of-00002.safetensors"),
38
+ os.path.join(tmpdir, "model-00002-of-00002.safetensors"),
39
+ ]
40
+
41
+ with pytest.raises(FileNotFoundError) as exc_info:
42
+ filter_duplicate_safetensors_files(
43
+ hf_weights_files=hf_weights_files,
44
+ hf_folder=tmpdir,
45
+ index_file="model.safetensors.index.json",
46
+ )
47
+
48
+ assert "model-00003-of-00002.safetensors" in str(exc_info.value)
49
+
50
+
51
+def test_filter_duplicate_safetensors_files_all_exist():
52
+ with tempfile.TemporaryDirectory() as tmpdir:
53
+ existing_files = []
54
+ for i in range(1, 3):
55
+ file_path = os.path.join(tmpdir, f"model-0000{i}-of-00002.safetensors")
56
+ with open(file_path, "wb") as f:
57
+ f.write(b"")
58
+ existing_files.append(file_path)
59
+
60
+ index_file = os.path.join(tmpdir, "model.safetensors.index.json")
61
+ index_content = {
62
+ "weight_map": {
63
+ "layer.0.weight": "model-00001-of-00002.safetensors",
64
+ "layer.1.weight": "model-00002-of-00002.safetensors",
65
+ }
66
+ }
67
+ with open(index_file, "w") as f:
68
+ json.dump(index_content, f)
69
+
70
+ filter_duplicate_safetensors_files(
71
+ hf_weights_files=existing_files,
72
+ hf_folder=tmpdir,
73
+ index_file="model.safetensors.index.json",
74
+ )
75
+
76
+
77
+if __name__ == "__main__":
78
+ test_filter_duplicate_safetensors_files_missing_weight()
79
+ test_filter_duplicate_safetensors_files_all_exist()
vllm/model_executor/model_loader/weight_utils.py
+8
@@ -595,6 +595,14 @@ def filter_duplicate_safetensors_files(
595
weight_files_in_index = set()
596
for weight_name in weight_map:
597
weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
598
+ # Check if files referenced in model.safetensors.index.json actually exist.
599
+ # Raise error if any file is missing.
600
+ hf_weights_files_set = set(hf_weights_files)
601
+ missing_files = weight_files_in_index - hf_weights_files_set
602
+ if missing_files:
603
+ raise FileNotFoundError(
604
+ f"Weight files referenced in index but missing: {missing_files}"
605
+ )
606
# Filter out any fields that are not found in the index file.
607
hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
608
return hf_weights_files