diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 6ac65b479..934579fd3 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -4118,7 +4118,11 @@ class RunaiModelStreamerLoader(BaseModelLoader): """Prepare weights for the model. If the model is not local, it will be downloaded.""" - from sglang.srt.utils.runai_utils import is_runai_obj_uri, list_safetensors + from sglang.srt.utils.runai_utils import ( + ObjectStorageModel, + is_runai_obj_uri, + list_safetensors, + ) is_object_storage_path = is_runai_obj_uri(model_name_or_path) if self._is_distributed is None: @@ -4160,6 +4164,10 @@ class RunaiModelStreamerLoader(BaseModelLoader): self.load_config.download_dir, revision, ) + if is_object_storage_path: + index_file = os.path.abspath( + os.path.join(ObjectStorageModel.get_path(hf_folder), index_file) + ) hf_weights_files = filter_duplicate_safetensors_files( hf_weights_files, hf_folder, index_file ) diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 55b5d6b3e..6bc867c35 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -830,7 +830,10 @@ def filter_duplicate_safetensors_files( if any(fnmatch.fnmatch(rel_path, pattern) for pattern in allow_patterns): files_to_validate.add(f) - missing_files = sorted(f for f in files_to_validate if not os.path.isfile(f)) + if "://" in hf_folder: + missing_files = sorted(files_to_validate.difference(hf_weights_files)) + else: + missing_files = sorted(f for f in files_to_validate if not os.path.isfile(f)) if missing_files: raise RuntimeError( f"{index_file} references {len(missing_files)} shard file(s) missing " @@ -876,7 +879,16 @@ def maybe_add_mtp_safetensors( # Check if mtp.safetensors exists and is not already in the file list mtp_path = os.path.join(hf_folder, "mtp.safetensors") - if not os.path.isfile(mtp_path) or mtp_path in hf_weights_files: + if mtp_path in hf_weights_files: + return hf_weights_files + + from sglang.srt.utils.runai_utils import is_runai_obj_uri, list_safetensors + + if is_runai_obj_uri(hf_folder): + mtp_exists = mtp_path in list_safetensors(hf_folder) + else: + mtp_exists = os.path.isfile(mtp_path) + if not mtp_exists: return hf_weights_files # mtp.safetensors exists but not in index - this is a bug diff --git a/test/registered/unit/model_loader/test_runai_model_streamer_loader.py b/test/registered/unit/model_loader/test_runai_model_streamer_loader.py index ec195a4f7..655445217 100644 --- a/test/registered/unit/model_loader/test_runai_model_streamer_loader.py +++ b/test/registered/unit/model_loader/test_runai_model_streamer_loader.py @@ -1,5 +1,8 @@ import concurrent.futures +import json +import os import sys +import tempfile import threading import unittest from types import SimpleNamespace @@ -20,6 +23,7 @@ from sglang.srt.models.deepseek_v4 import ( _dequant_fp8_wo_a_streaming, ) from sglang.srt.models.deepseek_v4_dspark import DeepseekV4ForCausalLMDSpark +from sglang.srt.utils import runai_utils from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -32,6 +36,80 @@ class _FakeModel: class TestRunaiModelStreamerLoader(CustomTestCase): + def test_remote_checkpoint_index(self): + for scheme in ("s3", "gs", "az"): + for relative_cache in (False, True): + with ( + self.subTest(scheme=scheme, relative_cache=relative_cache), + tempfile.TemporaryDirectory() as cache_dir, + ): + model_path = f"{scheme}://bucket/model" + shard = "model-00001-of-00001.safetensors" + files = [ + f"{model_path}/{shard}", + f"{model_path}/mtp.safetensors", + f"{model_path}/unused.safetensors", + ] + cache_root = ( + os.path.relpath(cache_dir) if relative_cache else cache_dir + ) + with ( + patch.object( + runai_utils.envs.SGLANG_CACHE_DIR, + "get", + return_value=cache_root, + ), + patch.object(loader_mod, "get_server_args", return_value=None), + patch.object( + runai_utils, "list_safetensors", return_value=files + ) as list_safetensors, + patch.object( + weight_utils, + "runai_safetensors_weights_iterator", + return_value=iter(()), + ) as streamer, + ): + metadata = runai_utils.ObjectStorageModel.get_path(model_path) + os.makedirs(metadata, exist_ok=True) + loader = loader_mod.RunaiModelStreamerLoader( + LoadConfig(load_format=LoadFormat.RUNAI_STREAMER) + ) + self.assertEqual( + loader._prepare_weights(model_path, revision=None), + (model_path, files), + ) + list_safetensors.assert_called_with(path=model_path) + + with open( + os.path.join(metadata, "model.safetensors.index.json"), "w" + ) as f: + json.dump({"weight_map": {"weight": shard}}, f) + self.assertEqual( + loader._prepare_weights(model_path, revision=None), + (model_path, [files[0]]), + ) + + loader.target_device_str = "cpu" + source = loader_mod.RunaiModelStreamerLoader.Source( + model_or_path=model_path, + revision=None, + model_config=cast( + ModelConfig, + SimpleNamespace( + hf_config=SimpleNamespace( + architectures=["Glm4MoeForCausalLMNextN"], + num_nextn_predict_layers=1, + ) + ), + ), + ) + list(loader._get_weights_iterator(source)) + streamer.assert_called_once_with(files[:2], True, "cpu") + + list_safetensors.return_value = [files[-1]] + with self.assertRaisesRegex(RuntimeError, shard): + loader._prepare_weights(model_path, revision=None) + def test_passes_quant_config_to_model_init(self): quant_config = object() fake_model = _FakeModel() diff --git a/test/registered/unit/model_loader/test_weight_utils.py b/test/registered/unit/model_loader/test_weight_utils.py index f63e642f0..9ebaa7203 100644 --- a/test/registered/unit/model_loader/test_weight_utils.py +++ b/test/registered/unit/model_loader/test_weight_utils.py @@ -4,8 +4,14 @@ import json import os import tempfile import unittest +from types import SimpleNamespace +from unittest.mock import patch -from sglang.srt.model_loader.weight_utils import filter_duplicate_safetensors_files +from sglang.srt.model_loader.weight_utils import ( + filter_duplicate_safetensors_files, + maybe_add_mtp_safetensors, +) +from sglang.srt.utils import runai_utils from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -85,17 +91,19 @@ class TestFilterDuplicateSafetensorsFiles(CustomTestCase): "vit": "vision_encoder/model.safetensors", }, ) - transformer = _touch( - self.folder, "transformer/diffusion_pytorch_model.safetensors" - ) + shard = "transformer/diffusion_pytorch_model.safetensors" + _touch(self.folder, shard) - result = filter_duplicate_safetensors_files( - hf_weights_files=[transformer], - hf_folder=self.folder, - index_file=INDEX_NAME, - allow_patterns=["transformer/*.safetensors"], - ) - self.assertEqual(result, [transformer]) + for folder in (self.folder, "s3://bucket/model"): + with self.subTest(folder=folder): + transformer = os.path.join(folder, shard) + result = filter_duplicate_safetensors_files( + hf_weights_files=[transformer], + hf_folder=folder, + index_file=os.path.join(self.folder, INDEX_NAME), + allow_patterns=["transformer/*.safetensors"], + ) + self.assertEqual(result, [transformer]) def test_missing_shard_inside_allow_patterns_raises(self): _write_index( @@ -106,18 +114,34 @@ class TestFilterDuplicateSafetensorsFiles(CustomTestCase): "vit": "vision_encoder/model.safetensors", }, ) - transformer = _touch( - self.folder, "transformer/model-00001-of-00002.safetensors" - ) + shard = "transformer/model-00001-of-00002.safetensors" + _touch(self.folder, shard) - with self.assertRaises(RuntimeError) as cm: - filter_duplicate_safetensors_files( - hf_weights_files=[transformer], - hf_folder=self.folder, - index_file=INDEX_NAME, - allow_patterns=["transformer/*.safetensors"], - ) - self.assertIn("model-00002-of-00002.safetensors", str(cm.exception)) + for folder in (self.folder, "s3://bucket/model"): + with ( + self.subTest(folder=folder), + self.assertRaisesRegex( + RuntimeError, r"model-00002-of-00002\.safetensors" + ), + ): + filter_duplicate_safetensors_files( + hf_weights_files=[os.path.join(folder, shard)], + hf_folder=folder, + index_file=os.path.join(self.folder, INDEX_NAME), + allow_patterns=["transformer/*.safetensors"], + ) + + def test_local_index_allows_subset_of_existing_shards(self): + _write_index( + self.folder, + {"w1": "first.safetensors", "w2": "second.safetensors"}, + ) + first = _touch(self.folder, "first.safetensors") + _touch(self.folder, "second.safetensors") + + result = filter_duplicate_safetensors_files([first], self.folder, INDEX_NAME) + + self.assertEqual(result, [first]) def test_single_file_model_no_index_returns_unchanged(self): # No index on disk (single-file / dummy / object-storage): early return. @@ -131,5 +155,49 @@ class TestFilterDuplicateSafetensorsFiles(CustomTestCase): self.assertEqual(result, [single]) +class TestMaybeAddMtpSafetensors(CustomTestCase): + def test_remote_mtp_guards(self): + folder = "s3://bucket/model" + model = f"{folder}/model.safetensors" + mtp = f"{folder}/mtp.safetensors" + cases = ( + ("Glm4MoeForCausalLM", 1, [model, mtp], [model, mtp], False), + ("Glm4MoeForCausalLM", 1, [model], [model], True), + ("Glm4MoeForCausalLM", 0, [model], [model, mtp], False), + ("LlamaForCausalLM", 1, [model], [model, mtp], False), + ) + for arch, nextn, selected, available, may_list in cases: + with ( + self.subTest(arch=arch, nextn=nextn, selected=selected), + patch.object( + runai_utils, "list_safetensors", return_value=available + ) as listing, + ): + config = SimpleNamespace( + architectures=[arch], num_nextn_predict_layers=nextn + ) + self.assertEqual( + maybe_add_mtp_safetensors(selected, folder, INDEX_NAME, config), + selected, + ) + if not may_list: + listing.assert_not_called() + + def test_local_unindexed_mtp(self): + with tempfile.TemporaryDirectory() as folder: + model = _touch(folder, "model.safetensors") + mtp = _touch(folder, "mtp.safetensors") + config = SimpleNamespace( + architectures=["Glm4MoeLiteForCausalLMNextN"], + num_nextn_predict_layers=1, + ) + with patch.object(runai_utils, "list_safetensors") as listing: + self.assertEqual( + maybe_add_mtp_safetensors([model], folder, INDEX_NAME, config), + [model, mtp], + ) + listing.assert_not_called() + + if __name__ == "__main__": unittest.main()