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>

tests/model_executor/model_loader/test_filter_duplicate_safetensors.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import os
+import tempfile
+
+import pytest
+
+from vllm.model_executor.model_loader.weight_utils import (
+    filter_duplicate_safetensors_files,
+)
+
+
+def test_filter_duplicate_safetensors_files_missing_weight():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        existing_file = os.path.join(tmpdir, "model-00001-of-00002.safetensors")
+        with open(existing_file, "wb") as f:
+            f.write(b"")
+
+        existing_file2 = os.path.join(tmpdir, "model-00002-of-00002.safetensors")
+        with open(existing_file2, "wb") as f:
+            f.write(b"")
+
+        index_file = os.path.join(tmpdir, "model.safetensors.index.json")
+        index_content = {
+            "weight_map": {
+                "layer.0.weight": "model-00001-of-00002.safetensors",
+                "layer.1.weight": "model-00002-of-00002.safetensors",
+                "layer.2.weight": "model-00003-of-00002.safetensors",
+            }
+        }
+        with open(index_file, "w") as f:
+            json.dump(index_content, f)
+
+        hf_weights_files = [
+            os.path.join(tmpdir, "model-00001-of-00002.safetensors"),
+            os.path.join(tmpdir, "model-00002-of-00002.safetensors"),
+        ]
+
+        with pytest.raises(FileNotFoundError) as exc_info:
+            filter_duplicate_safetensors_files(
+                hf_weights_files=hf_weights_files,
+                hf_folder=tmpdir,
+                index_file="model.safetensors.index.json",
+            )
+
+        assert "model-00003-of-00002.safetensors" in str(exc_info.value)
+
+
+def test_filter_duplicate_safetensors_files_all_exist():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        existing_files = []
+        for i in range(1, 3):
+            file_path = os.path.join(tmpdir, f"model-0000{i}-of-00002.safetensors")
+            with open(file_path, "wb") as f:
+                f.write(b"")
+            existing_files.append(file_path)
+
+        index_file = os.path.join(tmpdir, "model.safetensors.index.json")
+        index_content = {
+            "weight_map": {
+                "layer.0.weight": "model-00001-of-00002.safetensors",
+                "layer.1.weight": "model-00002-of-00002.safetensors",
+            }
+        }
+        with open(index_file, "w") as f:
+            json.dump(index_content, f)
+
+        filter_duplicate_safetensors_files(
+            hf_weights_files=existing_files,
+            hf_folder=tmpdir,
+            index_file="model.safetensors.index.json",
+        )
+
+
+if __name__ == "__main__":
+    test_filter_duplicate_safetensors_files_missing_weight()
+    test_filter_duplicate_safetensors_files_all_exist()

vllm/model_executor/model_loader/weight_utils.py

@@ -595,6 +595,14 @@ def filter_duplicate_safetensors_files(
     weight_files_in_index = set()
     for weight_name in weight_map:
         weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
+    # Check if files referenced in model.safetensors.index.json actually exist.
+    # Raise error if any file is missing.
+    hf_weights_files_set = set(hf_weights_files)
+    missing_files = weight_files_in_index - hf_weights_files_set
+    if missing_files:
+        raise FileNotFoundError(
+            f"Weight files referenced in index but missing: {missing_files}"
+        )
     # Filter out any fields that are not found in the index file.
     hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
     return hf_weights_files