[BugFix] Allow model_loader_extra_config with remote_instance + modelexpress backend (#34639)

Signed-off-by: joeltg <joel@reflection.ai>
Co-authored-by: mmangkad <176301910+mmangkad@users.noreply.github.com>

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:

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()