diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 613b823d6..dee6486b7 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -333,6 +333,28 @@ class BaseModelLoader(ABC): raise NotImplementedError +def _validate_default_loader_extra_config( + *, extra_config: dict, load_format: LoadFormat +) -> None: + allowed_keys = {"enable_multithread_load", "num_threads"} + if load_format == LoadFormat.FASTSAFETENSORS: + allowed_keys.add("enable_gds") + if "enable_gds" in extra_config and not isinstance( + extra_config["enable_gds"], bool + ): + raise ValueError( + "enable_gds in --model-loader-extra-config must be a boolean" + ) + + unexpected_keys = set(extra_config.keys()) - allowed_keys + if unexpected_keys: + raise ValueError( + f"Unexpected extra config keys for load format " + f"{load_format}: " + f"{unexpected_keys}" + ) + + class DefaultModelLoader(BaseModelLoader): """Model loader that can load different file types from disk.""" @@ -384,24 +406,10 @@ class DefaultModelLoader(BaseModelLoader): def __init__(self, load_config: LoadConfig): super().__init__(load_config) - extra_config = load_config.model_loader_extra_config - allowed_keys = {"enable_multithread_load", "num_threads"} - if load_config.load_format == LoadFormat.FASTSAFETENSORS: - allowed_keys.add("enable_gds") - if "enable_gds" in extra_config and not isinstance( - extra_config["enable_gds"], bool - ): - raise ValueError( - "enable_gds in --model-loader-extra-config must be a boolean" - ) - unexpected_keys = set(extra_config.keys()) - allowed_keys - - if unexpected_keys: - raise ValueError( - f"Unexpected extra config keys for load format " - f"{load_config.load_format}: " - f"{unexpected_keys}" - ) + _validate_default_loader_extra_config( + extra_config=load_config.model_loader_extra_config, + load_format=load_config.load_format, + ) def _maybe_download_from_modelscope( self, model: str, revision: Optional[str] @@ -3230,10 +3238,27 @@ class RemoteInstanceModelLoader(BaseModelLoader): def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: - raise ValueError( - f"Model loader extra config is not supported for " - f"load format {load_config.load_format}" - ) + if ( + load_config.remote_instance_weight_loader_backend + == RemoteInstanceWeightLoaderBackend.MODELEXPRESS + ): + # ModelExpress falls back to a DefaultModelLoader whenever no + # peer holds the weights, so it consumes this config; validate + # the keys here so a bad one fails on every rank instead of only + # on the ranks that end up taking the fallback. + _validate_default_loader_extra_config( + extra_config=load_config.model_loader_extra_config, + load_format=load_config.load_format, + ) + else: + # nccl and transfer_engine replace the native loader outright, + # so nothing would ever read the config. + raise ValueError( + f"Model loader extra config is not supported for " + f"load format {load_config.load_format} with " + f"remote instance weight loader backend " + f"{load_config.remote_instance_weight_loader_backend}" + ) self.remote_instance_transfer_engine_weight_info = None def download_model(self, model_config: ModelConfig) -> None: diff --git a/test/registered/unit/model_loader/test_remote_instance_loader.py b/test/registered/unit/model_loader/test_remote_instance_loader.py new file mode 100644 index 000000000..63f593e93 --- /dev/null +++ b/test/registered/unit/model_loader/test_remote_instance_loader.py @@ -0,0 +1,77 @@ +"""Unit tests for RemoteInstanceModelLoader construction - no server, no weights.""" + +import unittest + +import sglang.srt.model_loader.loader as loader_mod +from sglang.srt.configs.load_config import LoadConfig, LoadFormat +from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( + RemoteInstanceWeightLoaderBackend, +) +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") + +EXTRA_CONFIG = {"enable_multithread_load": True, "num_threads": 64} + +# ServerArgs forwards the backend as a plain str, so every case covers both that +# and the enum member; the loader's comparison holds only via the str mixin. +MODELEXPRESS_BACKENDS = ( + RemoteInstanceWeightLoaderBackend.MODELEXPRESS, + "modelexpress", +) +OTHER_BACKENDS = ( + RemoteInstanceWeightLoaderBackend.NCCL, + "nccl", + RemoteInstanceWeightLoaderBackend.TRANSFER_ENGINE, + "transfer_engine", +) +ALL_BACKENDS = MODELEXPRESS_BACKENDS + OTHER_BACKENDS + + +def _load_config(backend, extra_config=None): + return LoadConfig( + load_format=LoadFormat.REMOTE_INSTANCE, + model_loader_extra_config=extra_config or {}, + remote_instance_weight_loader_backend=backend, + ) + + +class TestRemoteInstanceModelLoaderExtraConfig(CustomTestCase): + def test_modelexpress_backend_accepts_extra_config(self): + for backend in MODELEXPRESS_BACKENDS: + with self.subTest(backend=backend): + loader = loader_mod.RemoteInstanceModelLoader( + _load_config(backend, EXTRA_CONFIG) + ) + # ModelExpress hands the extra config to the DefaultModelLoader + # it falls back to, so it must survive construction unmodified. + self.assertEqual( + loader.load_config.model_loader_extra_config, EXTRA_CONFIG + ) + + def test_modelexpress_backend_rejects_unknown_extra_config_key(self): + for backend in MODELEXPRESS_BACKENDS: + with self.subTest(backend=backend): + with self.assertRaises(ValueError): + loader_mod.RemoteInstanceModelLoader( + _load_config(backend, {"num_thread": 64}) + ) + + def test_other_backends_reject_extra_config(self): + for backend in OTHER_BACKENDS: + with self.subTest(backend=backend): + with self.assertRaises(ValueError): + loader_mod.RemoteInstanceModelLoader( + _load_config(backend, EXTRA_CONFIG) + ) + + def test_all_backends_construct_without_extra_config(self): + for backend in ALL_BACKENDS: + with self.subTest(backend=backend): + loader = loader_mod.RemoteInstanceModelLoader(_load_config(backend)) + self.assertFalse(loader.load_config.model_loader_extra_config) + + +if __name__ == "__main__": + unittest.main()