[Fix] Fail fast when a safetensors index references missing shard files (#32279)
This commit is contained in:
@@ -686,6 +686,16 @@ def filter_duplicate_safetensors_files(
|
|||||||
weight_files_in_index = set()
|
weight_files_in_index = set()
|
||||||
for weight_name in weight_map:
|
for weight_name in weight_map:
|
||||||
weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
|
weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
|
||||||
|
# Fail fast if the index references shard files that are not on disk (e.g. an
|
||||||
|
# incomplete or interrupted download). Otherwise those shards are silently
|
||||||
|
# dropped and the model loads with uninitialized weights.
|
||||||
|
missing_files = sorted(f for f in weight_files_in_index if not os.path.isfile(f))
|
||||||
|
if missing_files:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"{index_file} references {len(missing_files)} shard file(s) missing "
|
||||||
|
f"from {hf_folder} (incomplete download?): "
|
||||||
|
f"{[os.path.basename(f) for f in missing_files]}"
|
||||||
|
)
|
||||||
# Filter out any fields that are not found in the index file.
|
# 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]
|
hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
|
||||||
return hf_weights_files
|
return hf_weights_files
|
||||||
|
|||||||
@@ -0,0 +1,88 @@
|
|||||||
|
"""Unit tests for srt/model_loader/weight_utils.py shard-index consistency."""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
from sglang.srt.model_loader.weight_utils import filter_duplicate_safetensors_files
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
INDEX_NAME = "model.safetensors.index.json"
|
||||||
|
|
||||||
|
|
||||||
|
def _write_index(folder, weight_map):
|
||||||
|
with open(os.path.join(folder, INDEX_NAME), "w") as f:
|
||||||
|
json.dump({"weight_map": weight_map}, f)
|
||||||
|
|
||||||
|
|
||||||
|
def _touch(folder, name):
|
||||||
|
path = os.path.join(folder, name)
|
||||||
|
open(path, "w").close()
|
||||||
|
return path
|
||||||
|
|
||||||
|
|
||||||
|
class TestFilterDuplicateSafetensorsFiles(CustomTestCase):
|
||||||
|
def setUp(self):
|
||||||
|
self._tmp = tempfile.TemporaryDirectory()
|
||||||
|
self.folder = self._tmp.name
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
self._tmp.cleanup()
|
||||||
|
|
||||||
|
def test_missing_shard_raises(self):
|
||||||
|
# Index lists two shards, only one on disk (interrupted download).
|
||||||
|
_write_index(
|
||||||
|
self.folder,
|
||||||
|
{
|
||||||
|
"w1": "model-00001-of-00002.safetensors",
|
||||||
|
"w2": "model-00002-of-00002.safetensors",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
present = _touch(self.folder, "model-00001-of-00002.safetensors")
|
||||||
|
|
||||||
|
with self.assertRaises(RuntimeError) as cm:
|
||||||
|
filter_duplicate_safetensors_files(
|
||||||
|
hf_weights_files=[present],
|
||||||
|
hf_folder=self.folder,
|
||||||
|
index_file=INDEX_NAME,
|
||||||
|
)
|
||||||
|
self.assertIn("model-00002-of-00002.safetensors", str(cm.exception))
|
||||||
|
|
||||||
|
def test_complete_checkpoint_filters_non_indexed(self):
|
||||||
|
# All indexed shards present; a non-indexed duplicate is still filtered out.
|
||||||
|
_write_index(
|
||||||
|
self.folder,
|
||||||
|
{
|
||||||
|
"w1": "model-00001-of-00002.safetensors",
|
||||||
|
"w2": "model-00002-of-00002.safetensors",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
shard1 = _touch(self.folder, "model-00001-of-00002.safetensors")
|
||||||
|
shard2 = _touch(self.folder, "model-00002-of-00002.safetensors")
|
||||||
|
extra = _touch(self.folder, "consolidated.safetensors")
|
||||||
|
|
||||||
|
result = filter_duplicate_safetensors_files(
|
||||||
|
hf_weights_files=[shard1, shard2, extra],
|
||||||
|
hf_folder=self.folder,
|
||||||
|
index_file=INDEX_NAME,
|
||||||
|
)
|
||||||
|
self.assertEqual(sorted(result), sorted([shard1, shard2]))
|
||||||
|
|
||||||
|
def test_single_file_model_no_index_returns_unchanged(self):
|
||||||
|
# No index on disk (single-file / dummy / object-storage): early return.
|
||||||
|
single = _touch(self.folder, "model.safetensors")
|
||||||
|
|
||||||
|
result = filter_duplicate_safetensors_files(
|
||||||
|
hf_weights_files=[single],
|
||||||
|
hf_folder=self.folder,
|
||||||
|
index_file=INDEX_NAME,
|
||||||
|
)
|
||||||
|
self.assertEqual(result, [single])
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user