From 8d03a8172abc7f64b6f873276929457a0c7729b6 Mon Sep 17 00:00:00 2001 From: YitaoYuan Date: Sun, 20 Sep 2026 18:52:03 +0800 Subject: [PATCH 1/8] [Model Loader] Add InstantTensor loader support --- docs/docs/advanced_features/model_loading.mdx | 22 +++ .../advanced_features/server_arguments.mdx | 6 +- python/pyproject.toml | 2 + python/sglang/srt/arg_groups/choices.py | 1 + python/sglang/srt/arg_groups/fields/model.py | 2 + python/sglang/srt/configs/load_config.py | 3 + python/sglang/srt/model_loader/loader.py | 10 +- .../sglang/srt/model_loader/weight_utils.py | 38 +++++ python/sglang/test/runners.py | 2 + .../test_instanttensor_loader.py | 48 +++++++ .../model_loader/test_instanttensor_loader.py | 134 ++++++++++++++++++ 11 files changed, 262 insertions(+), 6 deletions(-) create mode 100644 test/registered/model_loading/test_instanttensor_loader.py create mode 100644 test/registered/unit/model_loader/test_instanttensor_loader.py diff --git a/docs/docs/advanced_features/model_loading.mdx b/docs/docs/advanced_features/model_loading.mdx index eb67e3c4b09b..99502b469e56 100644 --- a/docs/docs/advanced_features/model_loading.mdx +++ b/docs/docs/advanced_features/model_loading.mdx @@ -46,6 +46,10 @@ Set with `--load-format`: safetensors Load weights in the safetensors format. + + instanttensor + Load Safetensors weights quickly with InstantTensor's distributed loading and pipelined prefetching. + pt Load weights in the PyTorch .bin format. @@ -101,6 +105,24 @@ Set with `--load-format`: +## Load with InstantTensor + +InstantTensor provides high-performance loading for Safetensors checkpoints through distributed loading and pipelined prefetching. Install the optional dependency: + +```bash +pip install "sglang[instanttensor]" +``` + +Enable it explicitly when you launch the server: + +```bash +python -m sglang.launch_server \ + --model-path Qwen/Qwen3.6-35B-A3B \ + --load-format instanttensor +``` + +For tensor-parallel and multi-node deployments, InstantTensor uses SGLang's initialized distributed process group. The standard `auto` and `safetensors` paths remain unchanged. + ## Model loader extra config `--model-loader-extra-config` takes a JSON string passed to the loader selected by `--load-format`. diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index c386e6d30154..5be4c7d668f7 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -144,9 +144,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--load-format` - The format of the model weights to load. "auto" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. "pt" will load the weights in the pytorch bin format. "safetensors" will load the weights in the safetensors format. "npcache" will load the weights in pytorch format and store a numpy cache to speed up the loading. "dummy" will initialize the weights with random values, which is mainly for profiling."gguf" will load the weights in the gguf format. "bitsandbytes" will load the weights using bitsandbytes quantization."layered" loads weights layer by layer so that one can quantize a layer before loading another to make the peak memory envelope smaller. "flash_rl" will load the weights in flash_rl format. "fastsafetensors" and "private" are also supported. "runai_streamer" enables direct model loading from object storage and shared file systems. + The format of the model weights to load. "auto" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. "pt" will load the weights in the pytorch bin format. "safetensors" will load the weights in the safetensors format. "instanttensor" will load Safetensors weights using InstantTensor's high-performance distributed loader. "npcache" will load the weights in pytorch format and store a numpy cache to speed up the loading. "dummy" will initialize the weights with random values, which is mainly for profiling."gguf" will load the weights in the gguf format. "bitsandbytes" will load the weights using bitsandbytes quantization."layered" loads weights layer by layer so that one can quantize a layer before loading another to make the peak memory envelope smaller. "flash_rl" will load the weights in flash_rl format. "fastsafetensors" and "private" are also supported. "runai_streamer" enables direct model loading from object storage and shared file systems. `auto` - auto, pt, safetensors, npcache, dummy, sharded_state, gguf, bitsandbytes, mistral, layered, flash_rl, remote, remote_instance, fastsafetensors, private, runai_streamer + auto, pt, safetensors, instanttensor, npcache, dummy, sharded_state, gguf, bitsandbytes, mistral, layered, flash_rl, remote, remote_instance, fastsafetensors, private, runai_streamer `--model-loader-extra-config` @@ -1562,7 +1562,7 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec `--speculative-draft-load-format` The format of the draft model weights to load. If not specified, will use the same format as `--load-format`. Use 'dummy' to initialize draft model weights with random values for profiling. `None` - auto, pt, safetensors, npcache, dummy, sharded_state, gguf, bitsandbytes, mistral, layered, flash_rl, remote, remote_instance, fastsafetensors, private, runai_streamer + auto, pt, safetensors, instanttensor, npcache, dummy, sharded_state, gguf, bitsandbytes, mistral, layered, flash_rl, remote, remote_instance, fastsafetensors, private, runai_streamer `--speculative-num-steps` diff --git a/python/pyproject.toml b/python/pyproject.toml index fbf52feed3ad..8a7ad652f4fc 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -108,6 +108,7 @@ default = true [project.optional-dependencies] checkpoint-engine = ["checkpoint-engine==0.1.2"] +instanttensor = ["instanttensor>=0.1.5"] runai = ["runai-model-streamer[s3,gcs,azure]>=0.15.7"] diffusion = [ "addict==2.4.0", @@ -169,6 +170,7 @@ test = [ "diff-cover", "expecttest", "granian>=2.6.0", + "instanttensor>=0.1.5", "jsonlines", "lm-eval[api]>=0.4.9.2", "matplotlib", diff --git a/python/sglang/srt/arg_groups/choices.py b/python/sglang/srt/arg_groups/choices.py index 1ac5b914cb7e..6d0874722da2 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/python/sglang/srt/arg_groups/choices.py @@ -10,6 +10,7 @@ "auto", "pt", "safetensors", + "instanttensor", "npcache", "dummy", "sharded_state", diff --git a/python/sglang/srt/arg_groups/fields/model.py b/python/sglang/srt/arg_groups/fields/model.py index 606279c01a8a..9c1415dd16f8 100644 --- a/python/sglang/srt/arg_groups/fields/model.py +++ b/python/sglang/srt/arg_groups/fields/model.py @@ -65,6 +65,8 @@ class Model(msgspec.Struct): "is not available. " '"pt" will load the weights in the pytorch bin format. ' '"safetensors" will load the weights in the safetensors format. ' + '"instanttensor" will load Safetensors weights using ' + "InstantTensor's high-performance distributed loader. " '"npcache" will load the weights in pytorch format and store ' "a numpy cache to speed up the loading. " '"dummy" will initialize the weights with random values, ' diff --git a/python/sglang/srt/configs/load_config.py b/python/sglang/srt/configs/load_config.py index 5a0e06d026dd..980f0162ff07 100644 --- a/python/sglang/srt/configs/load_config.py +++ b/python/sglang/srt/configs/load_config.py @@ -18,6 +18,7 @@ class LoadFormat(str, enum.Enum): AUTO = "auto" PT = "pt" SAFETENSORS = "safetensors" + INSTANTTENSOR = "instanttensor" NPCACHE = "npcache" DUMMY = "dummy" SHARDED_STATE = "sharded_state" @@ -50,6 +51,8 @@ class LoadConfig: not available. "pt" will load the weights in the pytorch bin format. "safetensors" will load the weights in the safetensors format. + "instanttensor" will load Safetensors weights using InstantTensor's + high-performance distributed loader. "npcache" will load the weights in pytorch format and store a numpy cache to speed up the loading. "dummy" will initialize the weights with random values, which is diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 71ecd9ac5592..e1fe78618e20 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -116,6 +116,7 @@ gguf_quant_weights_iterator, initialize_capture_safe_weights, initialize_dummy_weights, + instanttensor_weights_iterator, maybe_add_mtp_safetensors, multi_thread_pt_weights_iterator, np_cache_weights_iterator, @@ -473,9 +474,10 @@ def _prepare_weights( # Some quantized models use .pt files for storing the weights. if load_format == LoadFormat.AUTO: allow_patterns = ["*.safetensors", "*.bin"] - elif ( - load_format == LoadFormat.SAFETENSORS - or load_format == LoadFormat.FASTSAFETENSORS + elif load_format in ( + LoadFormat.SAFETENSORS, + LoadFormat.FASTSAFETENSORS, + LoadFormat.INSTANTTENSOR, ): use_safetensors = True allow_patterns = ["*.safetensors"] @@ -614,6 +616,8 @@ def _get_weights_iterator( hf_folder, hf_weights_files, ) + elif self.load_config.load_format == LoadFormat.INSTANTTENSOR: + weights_iterator = instanttensor_weights_iterator(hf_weights_files) elif use_safetensors: weight_loader_disable_mmap = get_model().weight_loader_disable_mmap configured_prefetch = get_model().weight_loader_prefetch_checkpoints diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 8e34d94f7d3c..84cf4b5a6172 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -57,6 +57,7 @@ ci_download_with_validation_and_retry, ci_validate_and_cleanup_local_snapshot, ) +from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import ( BAR_FORMAT, @@ -1163,6 +1164,43 @@ def safetensors_weights_iterator( _drop_file_cache_after_load(st_file) +def instanttensor_weights_iterator( + hf_weights_files: List[str], +) -> Generator[Tuple[str, torch.Tensor], None, None]: + """Iterate over Safetensors weights with InstantTensor.""" + try: + import instanttensor + except ImportError as e: + raise ImportError( + "Please install InstantTensor via `pip install instanttensor`." + ) from e + + distributed = torch.distributed.is_initialized() + if distributed: + world_group = get_parallel().world_group + process_group = world_group.device_group if world_group.world_size > 1 else None + else: + process_group = None + + device = current_platform.get_device(torch.cuda.current_device()) + enable_tqdm = not distributed or torch.distributed.get_rank() == 0 + with instanttensor.safe_open( + hf_weights_files, + framework="pt", + device=device, + process_group=process_group, + ) as f: + yield from tqdm( + f.tensors(), + total=len(f.keys()), + desc="Loading safetensors using InstantTensor", + disable=not enable_tqdm, + mininterval=1, + bar_format=BAR_FORMAT, + position=tqdm._get_free_pos(), + ) + + def fastsafetensors_weights_iterator( hf_weights_files: List[str], enable_gds: bool = True, diff --git a/python/sglang/test/runners.py b/python/sglang/test/runners.py index 363daf9ddcf0..20c1b37a7c60 100644 --- a/python/sglang/test/runners.py +++ b/python/sglang/test/runners.py @@ -604,6 +604,7 @@ def __init__( lora_eviction_policy: str = "lru", enable_deterministic_inference: bool = False, lora_drain_wait_threshold: float = 0.0, + load_format: str = "auto", ): self.model_type = model_type self.is_generation = model_type == "generation" @@ -630,6 +631,7 @@ def __init__( dtype=get_dtype_str(torch_dtype), port=port, model_impl=model_impl, + load_format=load_format, mem_fraction_static=mem_fraction_static, trust_remote_code=trust_remote_code, is_embedding=not self.is_generation, diff --git a/test/registered/model_loading/test_instanttensor_loader.py b/test/registered/model_loading/test_instanttensor_loader.py new file mode 100644 index 000000000000..58233b046b13 --- /dev/null +++ b/test/registered/model_loading/test_instanttensor_loader.py @@ -0,0 +1,48 @@ +import multiprocessing as mp +import os +import unittest + +import torch + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.runners import SRTRunner, check_close_model_outputs +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=80, stage="weekly", runner_config="1-gpu-small") + +MODEL = os.getenv("SGLANG_INSTANTTENSOR_TEST_MODEL", "Qwen/Qwen2-0.5B") + + +class TestInstantTensorLoader(CustomTestCase): + @classmethod + def setUpClass(cls): + mp.set_start_method("spawn", force=True) + + def test_generation_matches_safetensors(self): + prompts = ["The capital of France is"] + runner_args = dict( + torch_dtype=torch.float16, + model_type="generation", + disable_cuda_graph=True, + disable_radix_cache=True, + max_total_tokens=256, + ) + + with SRTRunner(MODEL, load_format="safetensors", **runner_args) as runner: + expected = runner.forward(prompts, max_new_tokens=16) + + with SRTRunner(MODEL, load_format="instanttensor", **runner_args) as runner: + actual = runner.forward(prompts, max_new_tokens=16) + + check_close_model_outputs( + hf_outputs=expected, + srt_outputs=actual, + prefill_tolerance=1e-6, + decode_tolerance=1e-6, + rouge_l_tolerance=1.0, + debug_text="safetensors vs InstantTensor weight loading", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_loader/test_instanttensor_loader.py b/test/registered/unit/model_loader/test_instanttensor_loader.py new file mode 100644 index 000000000000..4d4e7741a78f --- /dev/null +++ b/test/registered/unit/model_loader/test_instanttensor_loader.py @@ -0,0 +1,134 @@ +import sys +import tempfile +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +import sglang.srt.model_loader.loader as loader_mod +import sglang.srt.model_loader.weight_utils as weight_utils +from sglang.srt.configs.load_config import LoadConfig, LoadFormat +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _FakeSafeOpen: + def __init__(self, tensors): + self._tensors = tensors + + def __enter__(self): + return self + + def __exit__(self, *_args): + pass + + def keys(self): + return [name for name, _ in self._tensors] + + def tensors(self): + yield from self._tensors + + +class TestInstantTensorLoader(CustomTestCase): + def test_iterator_uses_current_device_and_world_group(self): + tensors = [("weight", torch.tensor([1]))] + device_group = object() + calls = [] + + def safe_open(files, **kwargs): + calls.append((files, kwargs)) + return _FakeSafeOpen(tensors) + + files = ["b", "a"] + module = SimpleNamespace(safe_open=safe_open) + with ( + patch.dict(sys.modules, {"instanttensor": module}), + patch.object( + weight_utils.current_platform, + "get_device", + return_value=torch.device("cuda:0"), + ), + patch.object(weight_utils.torch.cuda, "current_device", return_value=0), + patch.object( + weight_utils, + "get_parallel", + return_value=SimpleNamespace( + world_group=SimpleNamespace( + world_size=2, + device_group=device_group, + ) + ), + ), + patch.object( + weight_utils.torch.distributed, + "is_initialized", + return_value=True, + ), + patch.object(weight_utils.torch.distributed, "get_rank", return_value=0), + ): + result = list(weight_utils.instanttensor_weights_iterator(files)) + + self.assertEqual(result, tensors) + self.assertIs(calls[0][0], files) + self.assertEqual(calls[0][1]["framework"], "pt") + self.assertEqual(calls[0][1]["device"], torch.device("cuda:0")) + self.assertIs(calls[0][1]["process_group"], device_group) + + def test_iterator_without_initialized_world_group(self): + module = SimpleNamespace( + safe_open=lambda files, **kwargs: _FakeSafeOpen([]), + ) + with ( + patch.dict(sys.modules, {"instanttensor": module}), + patch.object( + weight_utils.current_platform, + "get_device", + return_value=torch.device("cuda:0"), + ), + patch.object(weight_utils.torch.cuda, "current_device", return_value=0), + patch.object(weight_utils, "get_parallel") as get_parallel, + patch.object( + weight_utils.torch.distributed, + "is_initialized", + return_value=False, + ), + ): + self.assertEqual( + list( + weight_utils.instanttensor_weights_iterator(["model.safetensors"]) + ), + [], + ) + + get_parallel.assert_not_called() + + def test_iterator_reports_missing_dependency(self): + with ( + patch.dict(sys.modules, {"instanttensor": None}), + self.assertRaisesRegex(ImportError, "pip install instanttensor"), + ): + list(weight_utils.instanttensor_weights_iterator(["model.safetensors"])) + + def test_format_uses_safetensors_discovery(self): + with tempfile.TemporaryDirectory() as folder: + checkpoint = f"{folder}/model.safetensors" + open(checkpoint, "wb").close() + model_loader = loader_mod.DefaultModelLoader( + LoadConfig(load_format=LoadFormat.INSTANTTENSOR) + ) + with patch.object(loader_mod, "get_server_args", return_value=None): + _, weight_files, use_safetensors = model_loader._prepare_weights( + folder, + revision=None, + fall_back_to_pt=True, + ) + + self.assertEqual(weight_files, [checkpoint]) + self.assertTrue(use_safetensors) + + +if __name__ == "__main__": + unittest.main() From c301a202cf07d12b769c87401f1cbd944ef9bc3a Mon Sep 17 00:00:00 2001 From: YitaoYuan Date: Mon, 21 Sep 2026 21:37:09 +0800 Subject: [PATCH 2/8] [Model Loader] Expose InstantTensor tuning options via extra config --- docs/docs/advanced_features/model_loading.mdx | 27 ++++ python/pyproject.toml | 2 +- python/sglang/srt/model_loader/loader.py | 7 +- .../sglang/srt/model_loader/weight_utils.py | 37 ++++- python/sglang/test/runners.py | 2 + .../test_instanttensor_loader.py | 40 +++-- .../model_loader/test_instanttensor_loader.py | 139 +++++++++++++++++- 7 files changed, 238 insertions(+), 16 deletions(-) diff --git a/docs/docs/advanced_features/model_loading.mdx b/docs/docs/advanced_features/model_loading.mdx index 99502b469e56..13f6daf0ffa2 100644 --- a/docs/docs/advanced_features/model_loading.mdx +++ b/docs/docs/advanced_features/model_loading.mdx @@ -123,6 +123,33 @@ python -m sglang.launch_server \ For tensor-parallel and multi-node deployments, InstantTensor uses SGLang's initialized distributed process group. The standard `auto` and `safetensors` paths remain unchanged. +On NVIDIA CUDA, use `--model-loader-extra-config` to pass tuning options to +`instanttensor.safe_open` (requires InstantTensor 0.1.9 or later): + +```bash +python -m sglang.launch_server \ + --model-path Qwen/Qwen3.5-4B \ + --load-format instanttensor \ + --model-loader-extra-config '{"backend": "MMAP", "chunk_size": 4194304, "concurrency": 2, "io_depth": 2, "max_free_mem_usage": 0.25}' +``` + +| Option | Description | +| --- | --- | +| `buffer_size` | GPU tensor buffer size in bytes. | +| `chunk_size` | File I/O operation size in bytes. | +| `concurrency` | Worker threads for the `MMAP` and `CUFILE` backends. | +| `io_depth` | Maximum number of rank-local I/O operations in flight. | +| `max_free_mem_usage` | Maximum fraction of free GPU memory available to the tensor buffer. | +| `backend` | A backend name or an ordered list of candidates: `AIO`, `AIO_BUFFERED`, `URING`, `URING_BUFFERED`, `CUFILE`, `MMAP`, or the `BUFFERED` policy. | + +Explicit non-null values override the corresponding `INSTANTTENSOR_*` environment +variables. Omitted or `null` options retain InstantTensor's environment-variable +and automatic defaults. Backend names are case-sensitive; for example, +`{"backend": ["URING", "AIO"]}` selects the first available candidate. +SGLang manages `framework`, `device`, and `process_group`, and keeps InstantTensor's +default tensor-copy behavior. The generic `enable_multithread_load` and +`num_threads` options do not control InstantTensor; use `concurrency` instead. + ## Model loader extra config `--model-loader-extra-config` takes a JSON string passed to the loader selected by `--load-format`. diff --git a/python/pyproject.toml b/python/pyproject.toml index 8a7ad652f4fc..0f52001320e4 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -108,7 +108,7 @@ default = true [project.optional-dependencies] checkpoint-engine = ["checkpoint-engine==0.1.2"] -instanttensor = ["instanttensor>=0.1.5"] +instanttensor = ["instanttensor>=0.1.9"] runai = ["runai-model-streamer[s3,gcs,azure]>=0.15.7"] diffusion = [ "addict==2.4.0", diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index e1fe78618e20..650b39738266 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -103,6 +103,7 @@ ) from sglang.srt.environ import envs from sglang.srt.model_loader.weight_utils import ( + INSTANTTENSOR_CONFIG_KEYS, CheckpointFilePrefetchHandle, _prefetch_all_checkpoints, buffered_multi_thread_safetensors_weights_iterator, @@ -344,6 +345,8 @@ def _validate_default_loader_extra_config( *, extra_config: dict, load_format: LoadFormat ) -> None: allowed_keys = {"enable_multithread_load", "num_threads"} + if load_format == LoadFormat.INSTANTTENSOR: + allowed_keys.update(INSTANTTENSOR_CONFIG_KEYS) if load_format == LoadFormat.FASTSAFETENSORS: allowed_keys.add("enable_gds") if "enable_gds" in extra_config and not isinstance( @@ -617,7 +620,9 @@ def _get_weights_iterator( hf_weights_files, ) elif self.load_config.load_format == LoadFormat.INSTANTTENSOR: - weights_iterator = instanttensor_weights_iterator(hf_weights_files) + weights_iterator = instanttensor_weights_iterator( + hf_weights_files, extra_config=extra_config + ) elif use_safetensors: weight_loader_disable_mmap = get_model().weight_loader_disable_mmap configured_prefetch = get_model().weight_loader_prefetch_checkpoints diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 84cf4b5a6172..52c947f5eb5c 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -79,6 +79,17 @@ RUNAI_STREAMER_TENSOR_ATTR = "_sglang_runai_streamer_tensor" +INSTANTTENSOR_CONFIG_KEYS = frozenset( + { + "buffer_size", + "chunk_size", + "concurrency", + "io_depth", + "max_free_mem_usage", + "backend", + } +) + # Matches routed-expert weight keys in both HF-style layouts # (``...mlp.experts..{gate,up,down}_proj.weight``) and DeepSeek V4 @@ -1166,15 +1177,38 @@ def safetensors_weights_iterator( def instanttensor_weights_iterator( hf_weights_files: List[str], + extra_config: Optional[dict] = None, ) -> Generator[Tuple[str, torch.Tensor], None, None]: """Iterate over Safetensors weights with InstantTensor.""" try: import instanttensor except ImportError as e: raise ImportError( - "Please install InstantTensor via `pip install instanttensor`." + 'Please install InstantTensor via `pip install "instanttensor>=0.1.9"`.' ) from e + kwargs = { + key: value + for key, value in (extra_config or {}).items() + if key in INSTANTTENSOR_CONFIG_KEYS + } + backend = kwargs.get("backend") + if backend is not None: + names = [backend] if isinstance(backend, str) else backend + if not isinstance(names, list) or not names: + raise ValueError( + "InstantTensor backend must be a name or a non-empty list of names" + ) + available = { + **instanttensor.Backend.__members__, + **instanttensor.BackendPolicy.__members__, + } + if any(not isinstance(name, str) or name not in available for name in names): + raise ValueError( + f"Invalid InstantTensor backend {backend!r}; expected names from {sorted(available)}" + ) + kwargs["backend"] = [available[name] for name in names] + distributed = torch.distributed.is_initialized() if distributed: world_group = get_parallel().world_group @@ -1189,6 +1223,7 @@ def instanttensor_weights_iterator( framework="pt", device=device, process_group=process_group, + **kwargs, ) as f: yield from tqdm( f.tensors(), diff --git a/python/sglang/test/runners.py b/python/sglang/test/runners.py index 20c1b37a7c60..6d68f08a4096 100644 --- a/python/sglang/test/runners.py +++ b/python/sglang/test/runners.py @@ -605,6 +605,7 @@ def __init__( enable_deterministic_inference: bool = False, lora_drain_wait_threshold: float = 0.0, load_format: str = "auto", + model_loader_extra_config: Optional[dict[str, Any]] = None, ): self.model_type = model_type self.is_generation = model_type == "generation" @@ -632,6 +633,7 @@ def __init__( port=port, model_impl=model_impl, load_format=load_format, + model_loader_extra_config=json.dumps(model_loader_extra_config or {}), mem_fraction_static=mem_fraction_static, trust_remote_code=trust_remote_code, is_embedding=not self.is_generation, diff --git a/test/registered/model_loading/test_instanttensor_loader.py b/test/registered/model_loading/test_instanttensor_loader.py index 58233b046b13..9eddf20b0a9b 100644 --- a/test/registered/model_loading/test_instanttensor_loader.py +++ b/test/registered/model_loading/test_instanttensor_loader.py @@ -8,7 +8,7 @@ from sglang.test.runners import SRTRunner, check_close_model_outputs from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=80, stage="weekly", runner_config="1-gpu-small") +register_cuda_ci(est_time=120, stage="weekly", runner_config="1-gpu-small") MODEL = os.getenv("SGLANG_INSTANTTENSOR_TEST_MODEL", "Qwen/Qwen2-0.5B") @@ -31,17 +31,35 @@ def test_generation_matches_safetensors(self): with SRTRunner(MODEL, load_format="safetensors", **runner_args) as runner: expected = runner.forward(prompts, max_new_tokens=16) - with SRTRunner(MODEL, load_format="instanttensor", **runner_args) as runner: - actual = runner.forward(prompts, max_new_tokens=16) + configs = [ + {}, + { + "backend": "MMAP", + "chunk_size": 4 * 1024 * 1024, + "concurrency": 2, + "io_depth": 2, + "max_free_mem_usage": 0.25, + }, + ] + for extra_config in configs: + with self.subTest(extra_config=extra_config): + with SRTRunner( + MODEL, + load_format="instanttensor", + model_loader_extra_config=extra_config, + **runner_args, + ) as runner: + actual = runner.forward(prompts, max_new_tokens=16) - check_close_model_outputs( - hf_outputs=expected, - srt_outputs=actual, - prefill_tolerance=1e-6, - decode_tolerance=1e-6, - rouge_l_tolerance=1.0, - debug_text="safetensors vs InstantTensor weight loading", - ) + self.assertEqual(expected.output_strs, actual.output_strs) + check_close_model_outputs( + hf_outputs=expected, + srt_outputs=actual, + prefill_tolerance=1e-6, + decode_tolerance=1e-6, + rouge_l_tolerance=1.0, + debug_text="safetensors vs InstantTensor weight loading", + ) if __name__ == "__main__": diff --git a/test/registered/unit/model_loader/test_instanttensor_loader.py b/test/registered/unit/model_loader/test_instanttensor_loader.py index 4d4e7741a78f..a0b1d2cb9857 100644 --- a/test/registered/unit/model_loader/test_instanttensor_loader.py +++ b/test/registered/unit/model_loader/test_instanttensor_loader.py @@ -1,8 +1,10 @@ +import json import sys import tempfile import unittest +from enum import Enum from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch import torch @@ -33,6 +35,138 @@ def tensors(self): class TestInstantTensorLoader(CustomTestCase): + def test_extra_config_reaches_safe_open(self): + class Backend(Enum): + MMAP = 0 + AIO = 1 + + class BackendPolicy(Enum): + BUFFERED = 0 + + options = { + "buffer_size": 1024 * 1024 * 1024, + "chunk_size": 4 * 1024 * 1024, + "concurrency": 2, + "io_depth": 3, + "max_free_mem_usage": 0.25, + "backend": ["BUFFERED", "MMAP"], + } + config = {**options, "enable_multithread_load": False, "num_threads": 4} + load_config = LoadConfig( + load_format="instanttensor", model_loader_extra_config=json.dumps(config) + ) + model_loader = loader_mod.DefaultModelLoader(load_config) + safe_open = MagicMock(return_value=_FakeSafeOpen([])) + module = SimpleNamespace( + safe_open=safe_open, Backend=Backend, BackendPolicy=BackendPolicy + ) + with ( + patch.dict(sys.modules, {"instanttensor": module}), + patch.object(weight_utils.torch.cuda, "current_device", return_value=0), + patch.object( + weight_utils.current_platform, + "get_device", + return_value=torch.device("cuda:0"), + ), + patch.object( + weight_utils.torch.distributed, "is_initialized", return_value=False + ), + patch.object( + model_loader, + "_prepare_weights", + return_value=("model", ["model.safetensors"], True), + ), + ): + self.assertEqual( + list( + model_loader._get_weights_iterator( + loader_mod.DefaultModelLoader.Source("model", None) + ) + ), + [], + ) + + safe_open.assert_called_once_with( + ["model.safetensors"], + framework="pt", + device=torch.device("cuda:0"), + process_group=None, + **{**options, "backend": [BackendPolicy.BUFFERED, Backend.MMAP]}, + ) + self.assertEqual(load_config.model_loader_extra_config, config) + + def test_backend_names(self): + class Backend(Enum): + MMAP = 0 + + class BackendPolicy(Enum): + BUFFERED = 0 + + safe_open = MagicMock(return_value=_FakeSafeOpen([])) + module = SimpleNamespace( + safe_open=safe_open, Backend=Backend, BackendPolicy=BackendPolicy + ) + with ( + patch.dict(sys.modules, {"instanttensor": module}), + patch.object(weight_utils.torch.cuda, "current_device", return_value=0), + patch.object( + weight_utils.current_platform, + "get_device", + return_value=torch.device("cuda:0"), + ), + patch.object( + weight_utils.torch.distributed, "is_initialized", return_value=False + ), + ): + for backend, expected in [("MMAP", [Backend.MMAP]), (None, None)]: + with self.subTest(backend=backend): + list( + weight_utils.instanttensor_weights_iterator( + [], {"backend": backend} + ) + ) + self.assertEqual(safe_open.call_args.kwargs["backend"], expected) + null_options = dict.fromkeys(weight_utils.INSTANTTENSOR_CONFIG_KEYS) + list(weight_utils.instanttensor_weights_iterator([], null_options)) + for key in null_options: + self.assertIsNone(safe_open.call_args.kwargs[key]) + for backend in ["unknown", [], ["MMAP", "unknown"], 1, [1], {"MMAP": True}]: + with ( + self.subTest(backend=backend), + self.assertRaisesRegex(ValueError, "backend"), + ): + list( + weight_utils.instanttensor_weights_iterator( + [], {"backend": backend} + ) + ) + + def test_extra_config_rejects_unknown_and_managed_options(self): + for key in ["typo", "device", "framework", "process_group", "copy", "load_now"]: + with ( + self.subTest(key=key), + self.assertRaisesRegex(ValueError, "Unexpected extra config"), + ): + loader_mod.DefaultModelLoader( + LoadConfig( + load_format="instanttensor", + model_loader_extra_config={key: None}, + ) + ) + + def test_instanttensor_options_do_not_change_other_formats(self): + for load_format in ["auto", "safetensors", "fastsafetensors"]: + with ( + self.subTest(load_format=load_format), + self.assertRaisesRegex(ValueError, "Unexpected extra config"), + ): + loader_mod.DefaultModelLoader( + LoadConfig( + load_format=load_format, + model_loader_extra_config={"io_depth": 2}, + ) + ) + def test_iterator_uses_current_device_and_world_group(self): tensors = [("weight", torch.tensor([1]))] device_group = object() @@ -76,6 +210,7 @@ def safe_open(files, **kwargs): self.assertEqual(calls[0][1]["framework"], "pt") self.assertEqual(calls[0][1]["device"], torch.device("cuda:0")) self.assertIs(calls[0][1]["process_group"], device_group) + self.assertEqual(set(calls[0][1]), {"framework", "device", "process_group"}) def test_iterator_without_initialized_world_group(self): module = SimpleNamespace( @@ -108,7 +243,7 @@ def test_iterator_without_initialized_world_group(self): def test_iterator_reports_missing_dependency(self): with ( patch.dict(sys.modules, {"instanttensor": None}), - self.assertRaisesRegex(ImportError, "pip install instanttensor"), + self.assertRaisesRegex(ImportError, 'pip install "instanttensor>=0.1.9"'), ): list(weight_utils.instanttensor_weights_iterator(["model.safetensors"])) From 4316a80ccca59ab5ab37dd3a4c70810c850adffd Mon Sep 17 00:00:00 2001 From: YitaoYuan Date: Wed, 23 Sep 2026 03:24:57 +0800 Subject: [PATCH 3/8] [Model Loader] Delegate InstantTensor option validation to the backend --- python/sglang/srt/model_loader/loader.py | 6 +- .../sglang/srt/model_loader/weight_utils.py | 18 +-- .../test_instanttensor_loader.py | 8 +- .../model_loader/test_instanttensor_loader.py | 124 +++++++++++++++--- 4 files changed, 111 insertions(+), 45 deletions(-) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 650b39738266..1a16ff8abd79 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -103,7 +103,6 @@ ) from sglang.srt.environ import envs from sglang.srt.model_loader.weight_utils import ( - INSTANTTENSOR_CONFIG_KEYS, CheckpointFilePrefetchHandle, _prefetch_all_checkpoints, buffered_multi_thread_safetensors_weights_iterator, @@ -344,9 +343,10 @@ def load_model( def _validate_default_loader_extra_config( *, extra_config: dict, load_format: LoadFormat ) -> None: - allowed_keys = {"enable_multithread_load", "num_threads"} if load_format == LoadFormat.INSTANTTENSOR: - allowed_keys.update(INSTANTTENSOR_CONFIG_KEYS) + return + + 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( diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 52c947f5eb5c..dab4b0afddc9 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -79,17 +79,6 @@ RUNAI_STREAMER_TENSOR_ATTR = "_sglang_runai_streamer_tensor" -INSTANTTENSOR_CONFIG_KEYS = frozenset( - { - "buffer_size", - "chunk_size", - "concurrency", - "io_depth", - "max_free_mem_usage", - "backend", - } -) - # Matches routed-expert weight keys in both HF-style layouts # (``...mlp.experts..{gate,up,down}_proj.weight``) and DeepSeek V4 @@ -1187,11 +1176,7 @@ def instanttensor_weights_iterator( 'Please install InstantTensor via `pip install "instanttensor>=0.1.9"`.' ) from e - kwargs = { - key: value - for key, value in (extra_config or {}).items() - if key in INSTANTTENSOR_CONFIG_KEYS - } + kwargs = dict(extra_config or {}) backend = kwargs.get("backend") if backend is not None: names = [backend] if isinstance(backend, str) else backend @@ -1223,6 +1208,7 @@ def instanttensor_weights_iterator( framework="pt", device=device, process_group=process_group, + copy=True, **kwargs, ) as f: yield from tqdm( diff --git a/test/registered/model_loading/test_instanttensor_loader.py b/test/registered/model_loading/test_instanttensor_loader.py index 9eddf20b0a9b..4fc70fc9a9b2 100644 --- a/test/registered/model_loading/test_instanttensor_loader.py +++ b/test/registered/model_loading/test_instanttensor_loader.py @@ -34,11 +34,9 @@ def test_generation_matches_safetensors(self): configs = [ {}, { - "backend": "MMAP", - "chunk_size": 4 * 1024 * 1024, - "concurrency": 2, - "io_depth": 2, - "max_free_mem_usage": 0.25, + "backend": "URING", + "chunk_size": 8 * 1024 * 1024, + "io_depth": 64, }, ] for extra_config in configs: diff --git a/test/registered/unit/model_loader/test_instanttensor_loader.py b/test/registered/unit/model_loader/test_instanttensor_loader.py index a0b1d2cb9857..6d712fa864d5 100644 --- a/test/registered/unit/model_loader/test_instanttensor_loader.py +++ b/test/registered/unit/model_loader/test_instanttensor_loader.py @@ -37,7 +37,7 @@ def tensors(self): class TestInstantTensorLoader(CustomTestCase): def test_extra_config_reaches_safe_open(self): class Backend(Enum): - MMAP = 0 + URING = 0 AIO = 1 class BackendPolicy(Enum): @@ -45,13 +45,11 @@ class BackendPolicy(Enum): options = { "buffer_size": 1024 * 1024 * 1024, - "chunk_size": 4 * 1024 * 1024, - "concurrency": 2, - "io_depth": 3, - "max_free_mem_usage": 0.25, - "backend": ["BUFFERED", "MMAP"], + "chunk_size": 8 * 1024 * 1024, + "io_depth": 64, + "backend": ["URING", "AIO"], } - config = {**options, "enable_multithread_load": False, "num_threads": 4} + config = options.copy() load_config = LoadConfig( load_format="instanttensor", model_loader_extra_config=json.dumps(config) ) @@ -91,7 +89,8 @@ class BackendPolicy(Enum): framework="pt", device=torch.device("cuda:0"), process_group=None, - **{**options, "backend": [BackendPolicy.BUFFERED, Backend.MMAP]}, + copy=True, + **{**options, "backend": [Backend.URING, Backend.AIO]}, ) self.assertEqual(load_config.model_loader_extra_config, config) @@ -118,7 +117,11 @@ class BackendPolicy(Enum): weight_utils.torch.distributed, "is_initialized", return_value=False ), ): - for backend, expected in [("MMAP", [Backend.MMAP]), (None, None)]: + for backend, expected in [ + ("MMAP", [Backend.MMAP]), + ("BUFFERED", [BackendPolicy.BUFFERED]), + (None, None), + ]: with self.subTest(backend=backend): list( weight_utils.instanttensor_weights_iterator( @@ -126,7 +129,14 @@ class BackendPolicy(Enum): ) ) self.assertEqual(safe_open.call_args.kwargs["backend"], expected) - null_options = dict.fromkeys(weight_utils.INSTANTTENSOR_CONFIG_KEYS) + null_options = dict.fromkeys( + [ + "backend", + "buffer_size", + "chunk_size", + "io_depth", + ] + ) list(weight_utils.instanttensor_weights_iterator([], null_options)) for key in null_options: self.assertIsNone(safe_open.call_args.kwargs[key]) @@ -141,18 +151,87 @@ class BackendPolicy(Enum): ) ) - def test_extra_config_rejects_unknown_and_managed_options(self): - for key in ["typo", "device", "framework", "process_group", "copy", "load_now"]: - with ( - self.subTest(key=key), - self.assertRaisesRegex(ValueError, "Unexpected extra config"), - ): - loader_mod.DefaultModelLoader( - LoadConfig( - load_format="instanttensor", - model_loader_extra_config={key: None}, + def test_duplicate_options_are_rejected_by_python(self): + def safe_open(filename, *, framework, device, process_group, copy): + self.fail("Duplicate arguments must fail before entering safe_open") + + with ( + patch.dict( + sys.modules, {"instanttensor": SimpleNamespace(safe_open=safe_open)} + ), + patch.object(weight_utils.torch.cuda, "current_device", return_value=0), + patch.object( + weight_utils.current_platform, + "get_device", + return_value=torch.device("cuda:0"), + ), + patch.object( + weight_utils.torch.distributed, "is_initialized", return_value=False + ), + ): + for key in ["filename", "device", "framework", "process_group", "copy"]: + with self.subTest(key=key): + model_loader = loader_mod.DefaultModelLoader( + LoadConfig( + load_format="instanttensor", + model_loader_extra_config={ + key: False if key == "copy" else None + }, + ) + ) + source = loader_mod.DefaultModelLoader.Source("model", None) + resolved = loader_mod.DefaultModelLoader.ResolvedSource( + source=source, + hf_folder="model", + weight_files=("model.safetensors",), + use_safetensors=True, ) + with self.assertRaisesRegex(TypeError, f"multiple values.*'{key}'"): + list( + model_loader._get_weights_iterator( + source, resolved_source=resolved + ) + ) + + def test_parameter_errors_propagate_from_instanttensor(self): + error = ValueError("chunk_size must be greater than zero") + + def safe_open(files, *, framework, device, process_group, copy, chunk_size): + self.assertEqual(chunk_size, 0) + raise error + + with ( + patch.dict( + sys.modules, {"instanttensor": SimpleNamespace(safe_open=safe_open)} + ), + patch.object(weight_utils.torch.cuda, "current_device", return_value=0), + patch.object( + weight_utils.current_platform, + "get_device", + return_value=torch.device("cuda:0"), + ), + patch.object( + weight_utils.torch.distributed, "is_initialized", return_value=False + ), + ): + model_loader = loader_mod.DefaultModelLoader( + LoadConfig( + load_format="instanttensor", + model_loader_extra_config={"chunk_size": 0}, ) + ) + source = loader_mod.DefaultModelLoader.Source("model", None) + resolved = loader_mod.DefaultModelLoader.ResolvedSource( + source=source, + hf_folder="model", + weight_files=("model.safetensors",), + use_safetensors=True, + ) + with self.assertRaises(ValueError) as raised: + list( + model_loader._get_weights_iterator(source, resolved_source=resolved) + ) + self.assertIs(raised.exception, error) def test_instanttensor_options_do_not_change_other_formats(self): for load_format in ["auto", "safetensors", "fastsafetensors"]: @@ -210,7 +289,10 @@ def safe_open(files, **kwargs): self.assertEqual(calls[0][1]["framework"], "pt") self.assertEqual(calls[0][1]["device"], torch.device("cuda:0")) self.assertIs(calls[0][1]["process_group"], device_group) - self.assertEqual(set(calls[0][1]), {"framework", "device", "process_group"}) + self.assertIs(calls[0][1]["copy"], True) + self.assertEqual( + set(calls[0][1]), {"framework", "device", "process_group", "copy"} + ) def test_iterator_without_initialized_world_group(self): module = SimpleNamespace( From 5bc26a41f500404427e240578c70f4be4444083f Mon Sep 17 00:00:00 2001 From: YitaoYuan Date: Wed, 23 Sep 2026 03:25:09 +0800 Subject: [PATCH 4/8] [Model Loader] Address InstantTensor documentation and dependency review --- docs/docs/advanced_features/model_loading.mdx | 69 +++++++------------ python/pyproject.toml | 2 +- 2 files changed, 25 insertions(+), 46 deletions(-) diff --git a/docs/docs/advanced_features/model_loading.mdx b/docs/docs/advanced_features/model_loading.mdx index 13f6daf0ffa2..7c81288a03e3 100644 --- a/docs/docs/advanced_features/model_loading.mdx +++ b/docs/docs/advanced_features/model_loading.mdx @@ -105,51 +105,6 @@ Set with `--load-format`: -## Load with InstantTensor - -InstantTensor provides high-performance loading for Safetensors checkpoints through distributed loading and pipelined prefetching. Install the optional dependency: - -```bash -pip install "sglang[instanttensor]" -``` - -Enable it explicitly when you launch the server: - -```bash -python -m sglang.launch_server \ - --model-path Qwen/Qwen3.6-35B-A3B \ - --load-format instanttensor -``` - -For tensor-parallel and multi-node deployments, InstantTensor uses SGLang's initialized distributed process group. The standard `auto` and `safetensors` paths remain unchanged. - -On NVIDIA CUDA, use `--model-loader-extra-config` to pass tuning options to -`instanttensor.safe_open` (requires InstantTensor 0.1.9 or later): - -```bash -python -m sglang.launch_server \ - --model-path Qwen/Qwen3.5-4B \ - --load-format instanttensor \ - --model-loader-extra-config '{"backend": "MMAP", "chunk_size": 4194304, "concurrency": 2, "io_depth": 2, "max_free_mem_usage": 0.25}' -``` - -| Option | Description | -| --- | --- | -| `buffer_size` | GPU tensor buffer size in bytes. | -| `chunk_size` | File I/O operation size in bytes. | -| `concurrency` | Worker threads for the `MMAP` and `CUFILE` backends. | -| `io_depth` | Maximum number of rank-local I/O operations in flight. | -| `max_free_mem_usage` | Maximum fraction of free GPU memory available to the tensor buffer. | -| `backend` | A backend name or an ordered list of candidates: `AIO`, `AIO_BUFFERED`, `URING`, `URING_BUFFERED`, `CUFILE`, `MMAP`, or the `BUFFERED` policy. | - -Explicit non-null values override the corresponding `INSTANTTENSOR_*` environment -variables. Omitted or `null` options retain InstantTensor's environment-variable -and automatic defaults. Backend names are case-sensitive; for example, -`{"backend": ["URING", "AIO"]}` selects the first available candidate. -SGLang manages `framework`, `device`, and `process_group`, and keeps InstantTensor's -default tensor-copy behavior. The generic `enable_multithread_load` and -`num_threads` options do not control InstantTensor; use `concurrency` instead. - ## Model loader extra config `--model-loader-extra-config` takes a JSON string passed to the loader selected by `--load-format`. @@ -200,6 +155,30 @@ python -m sglang.launch_server \ Use GPU Direct Storage. Set to false when the host does not provide the NVIDIA GPUDirect Storage kernel driver, such as in a gVisor sandbox. true + + instanttensor + buffer_size (int) + GPU tensor buffer size in bytes. See the InstantTensor GitHub homepage for the latest optional parameters. + null + + + instanttensor + chunk_size (int) + File I/O operation size in bytes. + null + + + instanttensor + io_depth (int) + Maximum number of rank-local I/O operations in flight. + null + + + instanttensor + backend (str / list) + A backend name or an ordered list of candidates: AIO, AIO_BUFFERED, URING, URING_BUFFERED, CUFILE, MMAP, or the BUFFERED policy. + null + bitsandbytes qlora_adapter_name_or_path (str) diff --git a/python/pyproject.toml b/python/pyproject.toml index 0f52001320e4..0d6b29fc701f 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -170,7 +170,7 @@ test = [ "diff-cover", "expecttest", "granian>=2.6.0", - "instanttensor>=0.1.5", + "instanttensor>=0.1.9", "jsonlines", "lm-eval[api]>=0.4.9.2", "matplotlib", From 5cdc9c9fbf9da9fbab0da709fe10473f82cde142 Mon Sep 17 00:00:00 2001 From: YitaoYuan Date: Fri, 25 Sep 2026 02:14:08 +0800 Subject: [PATCH 5/8] [Model Loader] Simplify InstantTensor tests and fix CI registration - Use a valid weekly single-GPU suite and isolate InstantTensor environment variables - Reduce E2E testing to Safetensors and one explicit MMAP configuration - Consolidate CPU tests around core integration behavior - Reject unsupported file suffixes at the iterator entry without changing checkpoint discovery --- python/sglang/srt/model_loader/loader.py | 2 + .../sglang/srt/model_loader/weight_utils.py | 7 + .../test_instanttensor_loader.py | 60 ++- .../model_loader/test_instanttensor_loader.py | 403 ++++++------------ 4 files changed, 157 insertions(+), 315 deletions(-) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 1a16ff8abd79..20e2ed19ee8b 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -344,6 +344,8 @@ def _validate_default_loader_extra_config( *, extra_config: dict, load_format: LoadFormat ) -> None: if load_format == LoadFormat.INSTANTTENSOR: + # Pass extra config directly to InstantTensor and let it report invalid + # options, so newer versions can add parameters without a SGLang update. return allowed_keys = {"enable_multithread_load", "num_threads"} diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index dab4b0afddc9..cd49fc356cf2 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -1169,6 +1169,13 @@ def instanttensor_weights_iterator( extra_config: Optional[dict] = None, ) -> Generator[Tuple[str, torch.Tensor], None, None]: """Iterate over Safetensors weights with InstantTensor.""" + unsupported_files = [f for f in hf_weights_files if not f.endswith(".safetensors")] + if unsupported_files: + raise ValueError( + "InstantTensor only supports .safetensors checkpoints; " + f"unsupported files: {unsupported_files}" + ) + try: import instanttensor except ImportError as e: diff --git a/test/registered/model_loading/test_instanttensor_loader.py b/test/registered/model_loading/test_instanttensor_loader.py index 4fc70fc9a9b2..6fcd805aba75 100644 --- a/test/registered/model_loading/test_instanttensor_loader.py +++ b/test/registered/model_loading/test_instanttensor_loader.py @@ -4,11 +4,12 @@ import torch +from sglang.srt.utils.common import temp_set_env from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.runners import SRTRunner, check_close_model_outputs from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=120, stage="weekly", runner_config="1-gpu-small") +register_cuda_ci(est_time=120, stage="weekly", runner_config="1-gpu-large") MODEL = os.getenv("SGLANG_INSTANTTENSOR_TEST_MODEL", "Qwen/Qwen2-0.5B") @@ -28,36 +29,33 @@ def test_generation_matches_safetensors(self): max_total_tokens=256, ) - with SRTRunner(MODEL, load_format="safetensors", **runner_args) as runner: - expected = runner.forward(prompts, max_new_tokens=16) - - configs = [ - {}, - { - "backend": "URING", - "chunk_size": 8 * 1024 * 1024, - "io_depth": 64, - }, - ] - for extra_config in configs: - with self.subTest(extra_config=extra_config): - with SRTRunner( - MODEL, - load_format="instanttensor", - model_loader_extra_config=extra_config, - **runner_args, - ) as runner: - actual = runner.forward(prompts, max_new_tokens=16) - - self.assertEqual(expected.output_strs, actual.output_strs) - check_close_model_outputs( - hf_outputs=expected, - srt_outputs=actual, - prefill_tolerance=1e-6, - decode_tolerance=1e-6, - rouge_l_tolerance=1.0, - debug_text="safetensors vs InstantTensor weight loading", - ) + with temp_set_env( + **{key: None for key in os.environ if key.startswith("INSTANTTENSOR_")} + ): + with SRTRunner(MODEL, load_format="safetensors", **runner_args) as runner: + expected = runner.forward(prompts, max_new_tokens=16) + + with SRTRunner( + MODEL, + load_format="instanttensor", + model_loader_extra_config={ + "backend": "MMAP", + "chunk_size": 8 * 1024 * 1024, + "io_depth": 2, + }, + **runner_args, + ) as runner: + actual = runner.forward(prompts, max_new_tokens=16) + + self.assertEqual(expected.output_strs, actual.output_strs) + check_close_model_outputs( + hf_outputs=expected, + srt_outputs=actual, + prefill_tolerance=1e-6, + decode_tolerance=1e-6, + rouge_l_tolerance=1.0, + debug_text="safetensors vs InstantTensor weight loading", + ) if __name__ == "__main__": diff --git a/test/registered/unit/model_loader/test_instanttensor_loader.py b/test/registered/unit/model_loader/test_instanttensor_loader.py index 6d712fa864d5..da8b03152d52 100644 --- a/test/registered/unit/model_loader/test_instanttensor_loader.py +++ b/test/registered/unit/model_loader/test_instanttensor_loader.py @@ -1,7 +1,7 @@ import json import sys -import tempfile import unittest +from contextlib import ExitStack from enum import Enum from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -10,13 +10,23 @@ import sglang.srt.model_loader.loader as loader_mod import sglang.srt.model_loader.weight_utils as weight_utils -from sglang.srt.configs.load_config import LoadConfig, LoadFormat +from sglang.srt.configs.load_config import LoadConfig from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=1, suite="base-a-test-cpu") +class _Backend(Enum): + MMAP = 0 + URING = 1 + AIO = 2 + + +class _BackendPolicy(Enum): + BUFFERED = 0 + + class _FakeSafeOpen: def __init__(self, tensors): self._tensors = tensors @@ -35,316 +45,141 @@ def tensors(self): class TestInstantTensorLoader(CustomTestCase): - def test_extra_config_reaches_safe_open(self): - class Backend(Enum): - URING = 0 - AIO = 1 - - class BackendPolicy(Enum): - BUFFERED = 0 + def setUp(self): + super().setUp() + patches = ExitStack() + self.addCleanup(patches.close) + self.tensors = [("weight", torch.tensor([1]))] + self.safe_open = MagicMock(return_value=_FakeSafeOpen(self.tensors)) + patches.enter_context( + patch.dict( + sys.modules, + { + "instanttensor": SimpleNamespace( + safe_open=self.safe_open, + Backend=_Backend, + BackendPolicy=_BackendPolicy, + ) + }, + ) + ) + patches.enter_context( + patch.object(weight_utils.torch.cuda, "current_device", return_value=1) + ) + self.get_device = patches.enter_context( + patch.object( + weight_utils.current_platform, + "get_device", + return_value=torch.device("cuda:1"), + ) + ) + self.is_initialized = patches.enter_context( + patch.object( + weight_utils.torch.distributed, "is_initialized", return_value=False + ) + ) + self.get_parallel = patches.enter_context( + patch.object(weight_utils, "get_parallel") + ) + def test_extra_config_reaches_safe_open(self): options = { "buffer_size": 1024 * 1024 * 1024, "chunk_size": 8 * 1024 * 1024, "io_depth": 64, "backend": ["URING", "AIO"], } - config = options.copy() - load_config = LoadConfig( - load_format="instanttensor", model_loader_extra_config=json.dumps(config) + config = LoadConfig( + load_format="instanttensor", model_loader_extra_config=json.dumps(options) ) - model_loader = loader_mod.DefaultModelLoader(load_config) - safe_open = MagicMock(return_value=_FakeSafeOpen([])) - module = SimpleNamespace( - safe_open=safe_open, Backend=Backend, BackendPolicy=BackendPolicy + model_loader = loader_mod.DefaultModelLoader(config) + source = loader_mod.DefaultModelLoader.Source("model", None) + resolved = loader_mod.DefaultModelLoader.ResolvedSource( + source=source, + hf_folder="model", + weight_files=("model.safetensors",), + use_safetensors=True, + ) + result = list( + model_loader._get_weights_iterator(source, resolved_source=resolved) ) - with ( - patch.dict(sys.modules, {"instanttensor": module}), - patch.object(weight_utils.torch.cuda, "current_device", return_value=0), - patch.object( - weight_utils.current_platform, - "get_device", - return_value=torch.device("cuda:0"), - ), - patch.object( - weight_utils.torch.distributed, "is_initialized", return_value=False - ), - patch.object( - model_loader, - "_prepare_weights", - return_value=("model", ["model.safetensors"], True), - ), - ): - self.assertEqual( - list( - model_loader._get_weights_iterator( - loader_mod.DefaultModelLoader.Source("model", None) - ) - ), - [], - ) - safe_open.assert_called_once_with( + self.assertEqual(result, self.tensors) + self.safe_open.assert_called_once_with( ["model.safetensors"], framework="pt", - device=torch.device("cuda:0"), + device=torch.device("cuda:1"), process_group=None, copy=True, - **{**options, "backend": [Backend.URING, Backend.AIO]}, - ) - self.assertEqual(load_config.model_loader_extra_config, config) - - def test_backend_names(self): - class Backend(Enum): - MMAP = 0 - - class BackendPolicy(Enum): - BUFFERED = 0 - - safe_open = MagicMock(return_value=_FakeSafeOpen([])) - module = SimpleNamespace( - safe_open=safe_open, Backend=Backend, BackendPolicy=BackendPolicy + **{**options, "backend": [_Backend.URING, _Backend.AIO]}, ) - with ( - patch.dict(sys.modules, {"instanttensor": module}), - patch.object(weight_utils.torch.cuda, "current_device", return_value=0), - patch.object( - weight_utils.current_platform, - "get_device", - return_value=torch.device("cuda:0"), - ), - patch.object( - weight_utils.torch.distributed, "is_initialized", return_value=False - ), - ): - for backend, expected in [ - ("MMAP", [Backend.MMAP]), - ("BUFFERED", [BackendPolicy.BUFFERED]), - (None, None), - ]: - with self.subTest(backend=backend): - list( - weight_utils.instanttensor_weights_iterator( - [], {"backend": backend} - ) + self.assertEqual(config.model_loader_extra_config, options) + self.get_device.assert_called_once_with(1) + self.get_parallel.assert_not_called() + + def test_backend_conversion(self): + for backend, expected in [ + ("MMAP", [_Backend.MMAP]), + ("BUFFERED", [_BackendPolicy.BUFFERED]), + (["BUFFERED", "MMAP"], [_BackendPolicy.BUFFERED, _Backend.MMAP]), + (None, None), + ]: + with self.subTest(backend=backend): + self.safe_open.reset_mock() + list( + weight_utils.instanttensor_weights_iterator( + [], {"backend": backend} ) - self.assertEqual(safe_open.call_args.kwargs["backend"], expected) - null_options = dict.fromkeys( - [ - "backend", - "buffer_size", - "chunk_size", - "io_depth", - ] - ) - list(weight_utils.instanttensor_weights_iterator([], null_options)) - for key in null_options: - self.assertIsNone(safe_open.call_args.kwargs[key]) - for backend in ["unknown", [], ["MMAP", "unknown"], 1, [1], {"MMAP": True}]: - with ( - self.subTest(backend=backend), - self.assertRaisesRegex(ValueError, "backend"), - ): + ) + self.assertEqual(self.safe_open.call_args.kwargs["backend"], expected) + self.safe_open.assert_called_once() + for backend in ["unknown", [], ["MMAP", "unknown"], 1]: + with self.subTest(backend=backend): + self.safe_open.reset_mock() + with self.assertRaisesRegex(ValueError, "backend"): list( weight_utils.instanttensor_weights_iterator( [], {"backend": backend} ) ) + self.safe_open.assert_not_called() - def test_duplicate_options_are_rejected_by_python(self): - def safe_open(filename, *, framework, device, process_group, copy): - self.fail("Duplicate arguments must fail before entering safe_open") - - with ( - patch.dict( - sys.modules, {"instanttensor": SimpleNamespace(safe_open=safe_open)} - ), - patch.object(weight_utils.torch.cuda, "current_device", return_value=0), - patch.object( - weight_utils.current_platform, - "get_device", - return_value=torch.device("cuda:0"), - ), - patch.object( - weight_utils.torch.distributed, "is_initialized", return_value=False - ), - ): - for key in ["filename", "device", "framework", "process_group", "copy"]: - with self.subTest(key=key): - model_loader = loader_mod.DefaultModelLoader( - LoadConfig( - load_format="instanttensor", - model_loader_extra_config={ - key: False if key == "copy" else None - }, + def test_iterator_uses_initialized_world_group(self): + self.is_initialized.return_value = True + device_group = object() + with patch.object(weight_utils.torch.distributed, "get_rank", return_value=0): + for world_size in (1, 2): + with self.subTest(world_size=world_size): + self.safe_open.reset_mock() + self.get_parallel.return_value = SimpleNamespace( + world_group=SimpleNamespace( + world_size=world_size, device_group=device_group ) ) - source = loader_mod.DefaultModelLoader.Source("model", None) - resolved = loader_mod.DefaultModelLoader.ResolvedSource( - source=source, - hf_folder="model", - weight_files=("model.safetensors",), - use_safetensors=True, - ) - with self.assertRaisesRegex(TypeError, f"multiple values.*'{key}'"): - list( - model_loader._get_weights_iterator( - source, resolved_source=resolved - ) + result = list( + weight_utils.instanttensor_weights_iterator( + ["model.safetensors"] ) - - def test_parameter_errors_propagate_from_instanttensor(self): - error = ValueError("chunk_size must be greater than zero") - - def safe_open(files, *, framework, device, process_group, copy, chunk_size): - self.assertEqual(chunk_size, 0) - raise error - - with ( - patch.dict( - sys.modules, {"instanttensor": SimpleNamespace(safe_open=safe_open)} - ), - patch.object(weight_utils.torch.cuda, "current_device", return_value=0), - patch.object( - weight_utils.current_platform, - "get_device", - return_value=torch.device("cuda:0"), - ), - patch.object( - weight_utils.torch.distributed, "is_initialized", return_value=False - ), - ): - model_loader = loader_mod.DefaultModelLoader( - LoadConfig( - load_format="instanttensor", - model_loader_extra_config={"chunk_size": 0}, - ) - ) - source = loader_mod.DefaultModelLoader.Source("model", None) - resolved = loader_mod.DefaultModelLoader.ResolvedSource( - source=source, - hf_folder="model", - weight_files=("model.safetensors",), - use_safetensors=True, - ) - with self.assertRaises(ValueError) as raised: - list( - model_loader._get_weights_iterator(source, resolved_source=resolved) - ) - self.assertIs(raised.exception, error) - - def test_instanttensor_options_do_not_change_other_formats(self): - for load_format in ["auto", "safetensors", "fastsafetensors"]: - with ( - self.subTest(load_format=load_format), - self.assertRaisesRegex(ValueError, "Unexpected extra config"), - ): - loader_mod.DefaultModelLoader( - LoadConfig( - load_format=load_format, - model_loader_extra_config={"io_depth": 2}, ) - ) - - def test_iterator_uses_current_device_and_world_group(self): - tensors = [("weight", torch.tensor([1]))] - device_group = object() - calls = [] - - def safe_open(files, **kwargs): - calls.append((files, kwargs)) - return _FakeSafeOpen(tensors) - - files = ["b", "a"] - module = SimpleNamespace(safe_open=safe_open) - with ( - patch.dict(sys.modules, {"instanttensor": module}), - patch.object( - weight_utils.current_platform, - "get_device", - return_value=torch.device("cuda:0"), - ), - patch.object(weight_utils.torch.cuda, "current_device", return_value=0), - patch.object( - weight_utils, - "get_parallel", - return_value=SimpleNamespace( - world_group=SimpleNamespace( - world_size=2, - device_group=device_group, + self.assertEqual(result, self.tensors) + self.safe_open.assert_called_once_with( + ["model.safetensors"], + framework="pt", + device=torch.device("cuda:1"), + process_group=device_group if world_size > 1 else None, + copy=True, ) - ), - ), - patch.object( - weight_utils.torch.distributed, - "is_initialized", - return_value=True, - ), - patch.object(weight_utils.torch.distributed, "get_rank", return_value=0), - ): - result = list(weight_utils.instanttensor_weights_iterator(files)) - - self.assertEqual(result, tensors) - self.assertIs(calls[0][0], files) - self.assertEqual(calls[0][1]["framework"], "pt") - self.assertEqual(calls[0][1]["device"], torch.device("cuda:0")) - self.assertIs(calls[0][1]["process_group"], device_group) - self.assertIs(calls[0][1]["copy"], True) - self.assertEqual( - set(calls[0][1]), {"framework", "device", "process_group", "copy"} - ) - - def test_iterator_without_initialized_world_group(self): - module = SimpleNamespace( - safe_open=lambda files, **kwargs: _FakeSafeOpen([]), - ) - with ( - patch.dict(sys.modules, {"instanttensor": module}), - patch.object( - weight_utils.current_platform, - "get_device", - return_value=torch.device("cuda:0"), - ), - patch.object(weight_utils.torch.cuda, "current_device", return_value=0), - patch.object(weight_utils, "get_parallel") as get_parallel, - patch.object( - weight_utils.torch.distributed, - "is_initialized", - return_value=False, - ), - ): - self.assertEqual( - list( - weight_utils.instanttensor_weights_iterator(["model.safetensors"]) - ), - [], - ) - - get_parallel.assert_not_called() - - def test_iterator_reports_missing_dependency(self): - with ( - patch.dict(sys.modules, {"instanttensor": None}), - self.assertRaisesRegex(ImportError, 'pip install "instanttensor>=0.1.9"'), - ): - list(weight_utils.instanttensor_weights_iterator(["model.safetensors"])) - - def test_format_uses_safetensors_discovery(self): - with tempfile.TemporaryDirectory() as folder: - checkpoint = f"{folder}/model.safetensors" - open(checkpoint, "wb").close() - model_loader = loader_mod.DefaultModelLoader( - LoadConfig(load_format=LoadFormat.INSTANTTENSOR) - ) - with patch.object(loader_mod, "get_server_args", return_value=None): - _, weight_files, use_safetensors = model_loader._prepare_weights( - folder, - revision=None, - fall_back_to_pt=True, - ) - self.assertEqual(weight_files, [checkpoint]) - self.assertTrue(use_safetensors) + def test_iterator_rejects_unsupported_files(self): + for files in (["model.pt"], ["model.bin"], ["model.safetensors", "model.pt"]): + with self.subTest(files=files): + with self.assertRaisesRegex( + ValueError, r"only supports \.safetensors" + ) as raised: + list(weight_utils.instanttensor_weights_iterator(files)) + self.assertIn(files[-1], str(raised.exception)) + self.safe_open.assert_not_called() + self.get_device.assert_not_called() if __name__ == "__main__": From e860e95a97e9bf57fae08fbafc71655da5239867 Mon Sep 17 00:00:00 2001 From: YitaoYuan Date: Mon, 28 Sep 2026 22:24:55 +0800 Subject: [PATCH 6/8] [Model Loader] Guard InstantTensor loading on CUDA-compatible devices Validate the selected device during server-argument checks and reject unsupported platforms before importing InstantTensor or accessing CUDA APIs. Allow CUDA and ROCm through their shared PyTorch cuda device interface, document the requirement, and cover the hardware guards with CPU unit tests. --- docs/docs/advanced_features/model_loading.mdx | 2 +- .../sglang/srt/arg_groups/validation_hook.py | 7 +++ .../sglang/srt/model_loader/weight_utils.py | 6 +++ .../model_loader/test_instanttensor_loader.py | 45 +++++++++++++++++++ .../unit/server_args/test_server_args.py | 31 +++++++++++++ 5 files changed, 90 insertions(+), 1 deletion(-) diff --git a/docs/docs/advanced_features/model_loading.mdx b/docs/docs/advanced_features/model_loading.mdx index 7c81288a03e3..dcbfdf8d6aba 100644 --- a/docs/docs/advanced_features/model_loading.mdx +++ b/docs/docs/advanced_features/model_loading.mdx @@ -48,7 +48,7 @@ Set with `--load-format`: instanttensor - Load Safetensors weights quickly with InstantTensor's distributed loading and pipelined prefetching. + Load Safetensors weights quickly with InstantTensor's distributed loading and pipelined prefetching. Supports CUDA and ROCm; requires PyTorch's cuda device interface. pt diff --git a/python/sglang/srt/arg_groups/validation_hook.py b/python/sglang/srt/arg_groups/validation_hook.py index 8a1e15c954a9..9ced3f6e8e28 100644 --- a/python/sglang/srt/arg_groups/validation_hook.py +++ b/python/sglang/srt/arg_groups/validation_hook.py @@ -262,6 +262,13 @@ def check_server_args(server_args: Any): "MindSpore model impl is only supported on Ascend npu." ) + # ROCm also uses PyTorch's "cuda" device type. + if cfg.load_format == "instanttensor" and cfg.device != "cuda": + raise ValueError( + "InstantTensor requires a CUDA-compatible device (including CUDA and ROCm); " + f"got {cfg.device!r}." + ) + # Check metrics labels if ( not cfg.tokenizer_metrics_custom_labels_header diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index cd49fc356cf2..7e12f7a3cdb1 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -1169,6 +1169,12 @@ def instanttensor_weights_iterator( extra_config: Optional[dict] = None, ) -> Generator[Tuple[str, torch.Tensor], None, None]: """Iterate over Safetensors weights with InstantTensor.""" + if current_platform.device_type != "cuda": + raise ValueError( + "InstantTensor requires a CUDA-compatible device (including CUDA and ROCm); " + f"got {current_platform.device_type!r}." + ) + unsupported_files = [f for f in hf_weights_files if not f.endswith(".safetensors")] if unsupported_files: raise ValueError( diff --git a/test/registered/unit/model_loader/test_instanttensor_loader.py b/test/registered/unit/model_loader/test_instanttensor_loader.py index da8b03152d52..d2d06493f745 100644 --- a/test/registered/unit/model_loader/test_instanttensor_loader.py +++ b/test/registered/unit/model_loader/test_instanttensor_loader.py @@ -11,6 +11,8 @@ import sglang.srt.model_loader.loader as loader_mod import sglang.srt.model_loader.weight_utils as weight_utils from sglang.srt.configs.load_config import LoadConfig +from sglang.srt.platforms.cuda import CudaSRTPlatform +from sglang.srt.platforms.rocm import RocmSRTPlatform from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -64,6 +66,9 @@ def setUp(self): ) ) patches.enter_context( + patch.object(weight_utils.current_platform, "device_type", "cuda") + ) + self.current_device = patches.enter_context( patch.object(weight_utils.torch.cuda, "current_device", return_value=1) ) self.get_device = patches.enter_context( @@ -117,6 +122,46 @@ def test_extra_config_reaches_safe_open(self): self.get_device.assert_called_once_with(1) self.get_parallel.assert_not_called() + def test_iterator_accepts_cuda_and_rocm(self): + for platform in (CudaSRTPlatform(), RocmSRTPlatform()): + with ( + self.subTest(platform=platform.device_name), + patch.object(weight_utils, "current_platform", platform), + ): + self.safe_open.reset_mock() + self.assertEqual( + list( + weight_utils.instanttensor_weights_iterator( + ["model.safetensors"] + ) + ), + self.tensors, + ) + self.safe_open.assert_called_once_with( + ["model.safetensors"], + framework="pt", + device=torch.device("cuda:1"), + process_group=None, + copy=True, + ) + + def test_iterator_rejects_non_cuda_devices_before_import(self): + for device_type in ("cpu", "xpu", "npu", "musa", "hpu", "mps"): + with ( + self.subTest(device_type=device_type), + patch.object(weight_utils.current_platform, "device_type", device_type), + patch.dict(sys.modules, {"instanttensor": None}), + self.assertRaisesRegex( + ValueError, "InstantTensor requires a CUDA-compatible device" + ) as raised, + ): + list(weight_utils.instanttensor_weights_iterator(["model.safetensors"])) + self.assertIn(repr(device_type), str(raised.exception)) + self.current_device.assert_not_called() + self.get_device.assert_not_called() + self.get_parallel.assert_not_called() + self.safe_open.assert_not_called() + def test_backend_conversion(self): for backend, expected in [ ("MMAP", [_Backend.MMAP]), diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 08a67e2e4e60..69d06527b23c 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -167,6 +167,37 @@ def test_weight_cache_daemon_allows_static_eplb(self): # daemon to build the same static EPLB layout as the engine. handle_load_format(args) + def test_instanttensor_requires_cuda_device(self): + for device in ("cpu", "xpu", "npu", "musa", "hpu", "mps"): + with ( + self.subTest(device=device), + self.assertRaisesRegex( + ValueError, "InstantTensor requires a CUDA-compatible device" + ), + ): + ServerArgs( + model_path="dummy", + served_model_name="dummy", + device=device, + load_format="instanttensor", + chunked_prefill_size=-1, + ).check_server_args() + + def test_instanttensor_device_check_allows_cuda_and_other_loaders(self): + for device, load_format in ( + ("cuda", "instanttensor"), + ("cpu", "auto"), + ("cpu", "safetensors"), + ): + with self.subTest(device=device, load_format=load_format): + ServerArgs( + model_path="dummy", + served_model_name="dummy", + device=device, + load_format=load_format, + chunked_prefill_size=-1, + ).check_server_args() + def test_enable_w4a4_mxfp4_megamoe_preserves_legacy_deepgemm_env(self): deepgemm_env = { "DG_USE_FP4_ACTS": "0", From 162b0056d83bd5233e9def83651b0982072dd40d Mon Sep 17 00:00:00 2001 From: YitaoYuan Date: Thu, 1 Oct 2026 19:32:30 +0800 Subject: [PATCH 7/8] Fix InstantTensor load groups for PP draft workers --- docs/docs/advanced_features/model_loading.mdx | 2 +- python/sglang/srt/configs/load_config.py | 18 ++++- python/sglang/srt/managers/tp_worker.py | 5 ++ .../sglang/srt/model_executor/model_runner.py | 9 ++- .../load_model_utils.py | 9 ++- python/sglang/srt/model_loader/loader.py | 4 +- .../sglang/srt/model_loader/weight_utils.py | 19 ++++-- .../sglang/srt/speculative/eagle_worker_v2.py | 9 +++ .../srt/speculative/pp_draft_embedding.py | 1 + .../test_instanttensor_loader.py | 5 +- .../model_loader/test_instanttensor_loader.py | 66 ++++++++++++++++++- 11 files changed, 133 insertions(+), 14 deletions(-) diff --git a/docs/docs/advanced_features/model_loading.mdx b/docs/docs/advanced_features/model_loading.mdx index dcbfdf8d6aba..64696b92a26c 100644 --- a/docs/docs/advanced_features/model_loading.mdx +++ b/docs/docs/advanced_features/model_loading.mdx @@ -48,7 +48,7 @@ Set with `--load-format`: instanttensor - Load Safetensors weights quickly with InstantTensor's distributed loading and pipelined prefetching. Supports CUDA and ROCm; requires PyTorch's cuda device interface. + Load Safetensors weights quickly with InstantTensor's distributed loading and pipelined prefetching. Supports CUDA and ROCm; requires PyTorch's cuda device interface. Distributed loading may retain communicator memory for the server lifetime. pt diff --git a/python/sglang/srt/configs/load_config.py b/python/sglang/srt/configs/load_config.py index 980f0162ff07..c27240d74e14 100644 --- a/python/sglang/srt/configs/load_config.py +++ b/python/sglang/srt/configs/load_config.py @@ -4,7 +4,7 @@ import enum import logging from dataclasses import dataclass, field -from typing import Any, List, Optional, Union +from typing import TYPE_CHECKING, Any, List, Optional, Union import orjson @@ -13,6 +13,17 @@ logger = logging.getLogger(__name__) +if TYPE_CHECKING: + from sglang.srt.distributed.parallel_state import GroupCoordinator + + +class _DefaultLoadGroup(enum.Enum): + DEFAULT = enum.auto() + + +_DEFAULT_LOAD_GROUP = _DefaultLoadGroup.DEFAULT +LoadGroup = Union["GroupCoordinator", _DefaultLoadGroup, None] + class LoadFormat(str, enum.Enum): AUTO = "auto" @@ -115,6 +126,11 @@ class LoadConfig: weight_cache_socket: Optional[str] = None # Path to daemon socket (for client mode) fallback_load_format: Union[str, "LoadFormat"] = LoadFormat.AUTO + # Runtime-only load group, currently used by InstantTensor; None disables + # collective loading. Other distributed loaders, such as fastsafetensors + # and runai_model_streamer, could also benefit from honoring this group. + load_group: LoadGroup = field(default=_DEFAULT_LOAD_GROUP, repr=False) + def __post_init__(self): model_loader_extra_config = self.model_loader_extra_config or {} if isinstance(model_loader_extra_config, str): diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index e5c3bc36a84d..55854b697bb5 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -22,6 +22,7 @@ import torch from sglang.srt.beam_search.logits_capture import capture_pre_sample_logits +from sglang.srt.configs.load_config import _DEFAULT_LOAD_GROUP, LoadGroup from sglang.srt.environ import envs from sglang.srt.managers.io_struct import ( DestroyWeightsUpdateGroupReqInput, @@ -291,9 +292,11 @@ def __init__( context_length: Optional[int] = None, draft_attention_backend: Optional[str] = None, random_seed: Optional[int] = None, + load_group: LoadGroup = _DEFAULT_LOAD_GROUP, ): # Parse args self.server_args = server_args + self.load_group = load_group self.gpu_id = gpu_id self.nccl_port = nccl_port self.is_draft_worker = is_draft_worker @@ -486,6 +489,7 @@ def _init_model_runner(self): memory_pool_config=self.memory_pool_config, draft_attention_backend=self.draft_attention_backend, draft_model_idx=0 if self.is_multi_layer_eagle else None, + load_group=self.load_group, ) def _init_multi_layer_eagle_model_runners(self): @@ -506,6 +510,7 @@ def _init_multi_layer_eagle_model_runners(self): memory_pool_config=self.memory_pool_config, draft_attention_backend=self.draft_attention_backend, draft_model_idx=i, + load_group=self.load_group, ) ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 937cd160aeb7..5fec3fb7d57a 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -25,7 +25,11 @@ import torch import torch.distributed as dist -from sglang.srt.configs.load_config import LoadConfig +from sglang.srt.configs.load_config import ( + _DEFAULT_LOAD_GROUP, + LoadConfig, + LoadGroup, +) from sglang.srt.configs.model_config import ( AttentionArch, ModelConfig, @@ -333,9 +337,11 @@ def __init__( memory_pool_config: Optional[MemoryPoolConfig] = None, draft_model_idx: Optional[int] = None, draft_attention_backend: Optional[str] = None, + load_group: LoadGroup = _DEFAULT_LOAD_GROUP, ): # Parse args self.mem_fraction_static = mem_fraction_static + self.load_group = load_group # Set on target by `_resolve_memory_pool_config`; passed in for draft # workers so they reuse target's resolved sizes (replaces legacy # `server_args._draft_pool_config` mutation hack). @@ -1208,6 +1214,7 @@ def load_model(self): draft_model_idx=self.draft_model_idx, weight_cache_mode=get_model().weight_cache_mode, weight_cache_socket=get_model().weight_cache_socket, + load_group=self.load_group, ) maybe_enable_ipc_weight_cache( diff --git a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py index c978df0eed26..1210e458dbaa 100644 --- a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py @@ -16,7 +16,12 @@ modelexpress_url_of, ) from sglang.srt.configs.device_config import DeviceConfig -from sglang.srt.configs.load_config import LoadConfig, LoadFormat +from sglang.srt.configs.load_config import ( + _DEFAULT_LOAD_GROUP, + LoadConfig, + LoadFormat, + LoadGroup, +) from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS from sglang.srt.debug_utils.tensor_dump_forward_hook import ( register_forward_hook_for_model, @@ -202,6 +207,7 @@ def build_load_config( draft_model_idx: int | None, weight_cache_mode: str, weight_cache_socket: str | None, + load_group: LoadGroup = _DEFAULT_LOAD_GROUP, ) -> LoadConfig: from sglang.srt.configs.modelopt_config import ModelOptConfig @@ -231,6 +237,7 @@ def build_load_config( draft_model_idx=draft_model_idx, weight_cache_mode=weight_cache_mode, weight_cache_socket=weight_cache_socket, + load_group=load_group, ) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 12dec3f47ff7..543dd90a7321 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -638,7 +638,9 @@ def _get_weights_iterator( ) elif self.load_config.load_format == LoadFormat.INSTANTTENSOR: weights_iterator = instanttensor_weights_iterator( - hf_weights_files, extra_config=extra_config + hf_weights_files, + extra_config=extra_config, + load_group=self.load_config.load_group, ) elif use_safetensors: weight_loader_disable_mmap = get_model().weight_loader_disable_mmap diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index f22c740dc306..8505ee084894 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -38,7 +38,11 @@ from pydantic import BaseModel, ConfigDict, ValidationInfo, model_validator from tqdm.auto import tqdm -from sglang.srt.configs.load_config import LoadConfig +from sglang.srt.configs.load_config import ( + _DEFAULT_LOAD_GROUP, + LoadConfig, + LoadGroup, +) from sglang.srt.configs.model_config import ( REQUANTIZATION_METHODS, ModelConfig, @@ -1170,6 +1174,7 @@ def safetensors_weights_iterator( def instanttensor_weights_iterator( hf_weights_files: List[str], extra_config: Optional[dict] = None, + load_group: LoadGroup = _DEFAULT_LOAD_GROUP, ) -> Generator[Tuple[str, torch.Tensor], None, None]: """Iterate over Safetensors weights with InstantTensor.""" if current_platform.device_type != "cuda": @@ -1211,11 +1216,13 @@ def instanttensor_weights_iterator( kwargs["backend"] = [available[name] for name in names] distributed = torch.distributed.is_initialized() - if distributed: - world_group = get_parallel().world_group - process_group = world_group.device_group if world_group.world_size > 1 else None - else: - process_group = None + if load_group is _DEFAULT_LOAD_GROUP: + load_group = get_parallel().world_group if distributed else None + process_group = ( + load_group.device_group + if load_group is not None and load_group.world_size > 1 + else None + ) device = current_platform.get_device(torch.cuda.current_device()) enable_tqdm = not distributed or torch.distributed.get_rank() == 0 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 738e98112ee0..8c66acda81a3 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -9,6 +9,7 @@ draft_topk1_argmax_only, draft_topk1_postprocess, ) +from sglang.srt.configs.load_config import _DEFAULT_LOAD_GROUP from sglang.srt.configs.model_config import get_dsa_mtp_topk_width from sglang.srt.environ import envs from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_runner import ( @@ -188,6 +189,13 @@ def __init__( self.draft_owns_attention = ( get_parallel().attn_dp_enabled and self.speculative_algorithm.is_eagle3() ) + # Only the last target PP stage loads the draft. Capture its TP group + # before draft_pp_context hides the target's pipeline topology. + load_group = ( + get_parallel().tp_group + if get_parallel().pp_size > 1 + else _DEFAULT_LOAD_GROUP + ) with ( draft_tp_context(self.draft_owns_attention), draft_pp_context(), @@ -203,6 +211,7 @@ def __init__( # The draft runs at absolute target positions. context_length=target_worker.model_runner.model_config.context_len, random_seed=target_worker.random_seed, + load_group=load_group, ) # Alias for better readability diff --git a/python/sglang/srt/speculative/pp_draft_embedding.py b/python/sglang/srt/speculative/pp_draft_embedding.py index f13a5125512a..3a335d111e5b 100644 --- a/python/sglang/srt/speculative/pp_draft_embedding.py +++ b/python/sglang/srt/speculative/pp_draft_embedding.py @@ -51,6 +51,7 @@ LoadFormat.AUTO, LoadFormat.SAFETENSORS, LoadFormat.FASTSAFETENSORS, + LoadFormat.INSTANTTENSOR, LoadFormat.MISTRAL, LoadFormat.PT, LoadFormat.NPCACHE, diff --git a/test/registered/model_loading/test_instanttensor_loader.py b/test/registered/model_loading/test_instanttensor_loader.py index 6fcd805aba75..916253e899c6 100644 --- a/test/registered/model_loading/test_instanttensor_loader.py +++ b/test/registered/model_loading/test_instanttensor_loader.py @@ -9,9 +9,9 @@ from sglang.test.runners import SRTRunner, check_close_model_outputs from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=120, stage="weekly", runner_config="1-gpu-large") +register_cuda_ci(est_time=120, stage="weekly", runner_config="2-gpu-large") -MODEL = os.getenv("SGLANG_INSTANTTENSOR_TEST_MODEL", "Qwen/Qwen2-0.5B") +MODEL = "Qwen/Qwen2-0.5B" class TestInstantTensorLoader(CustomTestCase): @@ -24,6 +24,7 @@ def test_generation_matches_safetensors(self): runner_args = dict( torch_dtype=torch.float16, model_type="generation", + tp_size=2, disable_cuda_graph=True, disable_radix_cache=True, max_total_tokens=256, diff --git a/test/registered/unit/model_loader/test_instanttensor_loader.py b/test/registered/unit/model_loader/test_instanttensor_loader.py index d2d06493f745..3f781e09d223 100644 --- a/test/registered/unit/model_loader/test_instanttensor_loader.py +++ b/test/registered/unit/model_loader/test_instanttensor_loader.py @@ -1,5 +1,6 @@ import json import sys +import tempfile import unittest from contextlib import ExitStack from enum import Enum @@ -7,6 +8,7 @@ from unittest.mock import MagicMock, patch import torch +from safetensors.torch import save_file import sglang.srt.model_loader.loader as loader_mod import sglang.srt.model_loader.weight_utils as weight_utils @@ -199,7 +201,8 @@ def test_iterator_uses_initialized_world_group(self): self.get_parallel.return_value = SimpleNamespace( world_group=SimpleNamespace( world_size=world_size, device_group=device_group - ) + ), + tp_group=SimpleNamespace(world_size=1, device_group=object()), ) result = list( weight_utils.instanttensor_weights_iterator( @@ -215,6 +218,44 @@ def test_iterator_uses_initialized_world_group(self): copy=True, ) + def test_explicit_load_group_reaches_safe_open(self): + self.is_initialized.return_value = True + world_group = SimpleNamespace(world_size=4, device_group=object()) + tp_group = SimpleNamespace(world_size=2, device_group=object()) + self.get_parallel.return_value = SimpleNamespace( + world_group=world_group, tp_group=tp_group + ) + + with patch.object(weight_utils.torch.distributed, "get_rank", return_value=2): + for group in ( + tp_group, + None, + SimpleNamespace(world_size=1, device_group=object()), + ): + with self.subTest(group=group): + self.safe_open.reset_mock() + config = LoadConfig(load_format="instanttensor", load_group=group) + loader = loader_mod.DefaultModelLoader(config) + source = loader.Source("model", None) + resolved = loader.ResolvedSource( + source=source, + hf_folder="model", + weight_files=("model.safetensors",), + use_safetensors=True, + ) + list(loader._get_weights_iterator(source, resolved_source=resolved)) + self.safe_open.assert_called_once() + self.assertIs( + self.safe_open.call_args.kwargs["process_group"], + tp_group.device_group if group is tp_group else None, + ) + self.get_parallel.assert_not_called() + list(weight_utils.instanttensor_weights_iterator(["model.safetensors"])) + self.assertIs( + self.safe_open.call_args.kwargs["process_group"], + world_group.device_group, + ) + def test_iterator_rejects_unsupported_files(self): for files in (["model.pt"], ["model.bin"], ["model.safetensors", "model.pt"]): with self.subTest(files=files): @@ -226,6 +267,29 @@ def test_iterator_rejects_unsupported_files(self): self.safe_open.assert_not_called() self.get_device.assert_not_called() + def test_pp_embedding_can_reopen_instanttensor_checkpoint(self): + from sglang.srt.speculative.pp_draft_embedding import ( + load_embedding_tensor, + prepare_checkpoint_files, + ) + + expected = torch.arange(8, dtype=torch.float32).reshape(2, 4) + name = "model.embed_tokens.weight" + with tempfile.TemporaryDirectory() as folder: + save_file({name: expected}, f"{folder}/model.safetensors") + with patch.object(loader_mod, "get_server_args", return_value=None): + hf_folder, weight_files, use_safetensors = prepare_checkpoint_files( + folder, + revision=None, + load_config=LoadConfig(load_format="instanttensor"), + ) + key, actual = load_embedding_tensor( + hf_folder, weight_files, use_safetensors=use_safetensors + ) + self.assertEqual(key, name) + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + self.safe_open.assert_not_called() + if __name__ == "__main__": unittest.main() From 67dec7648199c06f92b3232055a926498df50c5c Mon Sep 17 00:00:00 2001 From: YitaoYuan Date: Fri, 2 Oct 2026 01:10:59 +0800 Subject: [PATCH 8/8] [Model Loader] Fix InstantTensor groups for DSpark and PP disk updates --- .../model_runner_components/weight_updater.py | 10 ++- .../sglang/srt/model_loader/weight_utils.py | 9 +++ .../srt/speculative/draft_worker_common.py | 3 + .../dspark_components/dspark_worker_v2.py | 3 + .../sglang/srt/speculative/eagle_worker_v2.py | 8 +-- .../model_loader/test_instanttensor_loader.py | 62 ++++++++++--------- 6 files changed, 59 insertions(+), 36 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py index 5d852a613675..c9ec8f423b2a 100644 --- a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py +++ b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py @@ -14,7 +14,10 @@ post_load_weights, ) from sglang.srt.model_loader.utils import set_default_torch_dtype -from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + get_pp_stage_load_group, +) from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_model from sglang.srt.utils import ( @@ -184,7 +187,10 @@ def update_weights_from_disk( target_device = torch.device(self.device) self.model_config.model_path = model_path - load_config = LoadConfig(load_format=load_format) + # PP stages handle disk updates sequentially, unlike cold startup. + load_config = LoadConfig( + load_format=load_format, load_group=get_pp_stage_load_group() + ) # Only support DefaultModelLoader for now loader = get_model_loader(load_config, self.model_config) diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 8505ee084894..66cda9522a51 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -1171,6 +1171,15 @@ def safetensors_weights_iterator( prefetch_handle.stop() +def get_pp_stage_load_group() -> LoadGroup: + """Select a group for stage-local loads, not all-rank cold startup. + + Call before draft contexts override the target's parallel topology. + """ + parallel = get_parallel() + return parallel.tp_group if parallel.pp_size > 1 else _DEFAULT_LOAD_GROUP + + def instanttensor_weights_iterator( hf_weights_files: List[str], extra_config: Optional[dict] = None, diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index 9387550f48f5..8fa004b3db4e 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -6,6 +6,7 @@ import msgspec import torch +from sglang.srt.configs.load_config import _DEFAULT_LOAD_GROUP, LoadGroup from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode @@ -68,6 +69,7 @@ def build_draft_tp_worker( attention_backend_override: Optional[str] = None, draft_worker_cls: type[TpModelWorker] = TpModelWorker, random_seed: Optional[int] = None, + load_group: LoadGroup = _DEFAULT_LOAD_GROUP, ) -> DraftWorkerBundle: # An override names a draft-specific backend the caller has already # validated (e.g. a self-drafting architecture); it skips the generic @@ -89,6 +91,7 @@ def build_draft_tp_worker( nccl_port=nccl_port, is_draft_worker=True, random_seed=random_seed, + load_group=load_group, # The draft runs at absolute target positions. context_length=target_model_config.context_len, draft_attention_backend=draft_backend, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 6dad00afa39d..ca4dd763a5ec 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -20,6 +20,7 @@ PPProxyTensors, compute_position, ) +from sglang.srt.model_loader.weight_utils import get_pp_stage_load_group from sglang.srt.runtime_context import ( get_disagg, get_exec, @@ -185,6 +186,7 @@ def __init__( # Inside the draft scope the context answers the draft's narrowed rank. self._target_tp_rank = get_parallel().tp_rank + load_group = get_pp_stage_load_group() with draft_pp_context(), self._draft_context(): bundle = build_draft_tp_worker( server_args=server_args, @@ -197,6 +199,7 @@ def __init__( ), draft_worker_cls=draft_worker_cls, random_seed=target_worker.random_seed, + load_group=load_group, ) self._draft_worker = bundle.draft_worker self.draft_model_runner = bundle.draft_model_runner diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 8c66acda81a3..edcf45f34395 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -9,7 +9,6 @@ draft_topk1_argmax_only, draft_topk1_postprocess, ) -from sglang.srt.configs.load_config import _DEFAULT_LOAD_GROUP from sglang.srt.configs.model_config import get_dsa_mtp_topk_width from sglang.srt.environ import envs from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_runner import ( @@ -58,6 +57,7 @@ DecodeCudaGraphRunner, get_batch_sizes_to_capture, ) +from sglang.srt.model_loader.weight_utils import get_pp_stage_load_group from sglang.srt.runtime_context import ( get_context, get_device, @@ -191,11 +191,7 @@ def __init__( ) # Only the last target PP stage loads the draft. Capture its TP group # before draft_pp_context hides the target's pipeline topology. - load_group = ( - get_parallel().tp_group - if get_parallel().pp_size > 1 - else _DEFAULT_LOAD_GROUP - ) + load_group = get_pp_stage_load_group() with ( draft_tp_context(self.draft_owns_attention), draft_pp_context(), diff --git a/test/registered/unit/model_loader/test_instanttensor_loader.py b/test/registered/unit/model_loader/test_instanttensor_loader.py index 3f781e09d223..240c4cac521d 100644 --- a/test/registered/unit/model_loader/test_instanttensor_loader.py +++ b/test/registered/unit/model_loader/test_instanttensor_loader.py @@ -12,9 +12,10 @@ import sglang.srt.model_loader.loader as loader_mod import sglang.srt.model_loader.weight_utils as weight_utils -from sglang.srt.configs.load_config import LoadConfig +from sglang.srt.configs.load_config import _DEFAULT_LOAD_GROUP, LoadConfig from sglang.srt.platforms.cuda import CudaSRTPlatform from sglang.srt.platforms.rocm import RocmSRTPlatform +from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -89,7 +90,12 @@ def setUp(self): patch.object(weight_utils, "get_parallel") ) - def test_extra_config_reaches_safe_open(self): + def test_loader_passes_extra_config_and_load_group(self): + self.is_initialized.return_value = True + group = SimpleNamespace(world_size=2, device_group=object()) + self.get_parallel.return_value = SimpleNamespace( + world_group=SimpleNamespace(world_size=4, device_group=object()) + ) options = { "buffer_size": 1024 * 1024 * 1024, "chunk_size": 8 * 1024 * 1024, @@ -97,7 +103,9 @@ def test_extra_config_reaches_safe_open(self): "backend": ["URING", "AIO"], } config = LoadConfig( - load_format="instanttensor", model_loader_extra_config=json.dumps(options) + load_format="instanttensor", + model_loader_extra_config=json.dumps(options), + load_group=group, ) model_loader = loader_mod.DefaultModelLoader(config) source = loader_mod.DefaultModelLoader.Source("model", None) @@ -107,22 +115,21 @@ def test_extra_config_reaches_safe_open(self): weight_files=("model.safetensors",), use_safetensors=True, ) - result = list( - model_loader._get_weights_iterator(source, resolved_source=resolved) - ) + with patch.object(weight_utils.torch.distributed, "get_rank", return_value=0): + result = list( + model_loader._get_weights_iterator(source, resolved_source=resolved) + ) self.assertEqual(result, self.tensors) self.safe_open.assert_called_once_with( ["model.safetensors"], framework="pt", device=torch.device("cuda:1"), - process_group=None, + process_group=group.device_group, copy=True, **{**options, "backend": [_Backend.URING, _Backend.AIO]}, ) self.assertEqual(config.model_loader_extra_config, options) - self.get_device.assert_called_once_with(1) - self.get_parallel.assert_not_called() def test_iterator_accepts_cuda_and_rocm(self): for platform in (CudaSRTPlatform(), RocmSRTPlatform()): @@ -161,7 +168,6 @@ def test_iterator_rejects_non_cuda_devices_before_import(self): self.assertIn(repr(device_type), str(raised.exception)) self.current_device.assert_not_called() self.get_device.assert_not_called() - self.get_parallel.assert_not_called() self.safe_open.assert_not_called() def test_backend_conversion(self): @@ -218,7 +224,7 @@ def test_iterator_uses_initialized_world_group(self): copy=True, ) - def test_explicit_load_group_reaches_safe_open(self): + def test_iterator_uses_explicit_load_group(self): self.is_initialized.return_value = True world_group = SimpleNamespace(world_size=4, device_group=object()) tp_group = SimpleNamespace(world_size=2, device_group=object()) @@ -234,27 +240,28 @@ def test_explicit_load_group_reaches_safe_open(self): ): with self.subTest(group=group): self.safe_open.reset_mock() - config = LoadConfig(load_format="instanttensor", load_group=group) - loader = loader_mod.DefaultModelLoader(config) - source = loader.Source("model", None) - resolved = loader.ResolvedSource( - source=source, - hf_folder="model", - weight_files=("model.safetensors",), - use_safetensors=True, + list( + weight_utils.instanttensor_weights_iterator( + ["model.safetensors"], load_group=group + ) ) - list(loader._get_weights_iterator(source, resolved_source=resolved)) self.safe_open.assert_called_once() self.assertIs( self.safe_open.call_args.kwargs["process_group"], tp_group.device_group if group is tp_group else None, ) - self.get_parallel.assert_not_called() - list(weight_utils.instanttensor_weights_iterator(["model.safetensors"])) - self.assertIs( - self.safe_open.call_args.kwargs["process_group"], - world_group.device_group, - ) + + def test_pp_stage_load_group(self): + tp_group = object() + for pp_size in (1, 2): + with self.subTest(pp_size=pp_size): + self.get_parallel.return_value = SimpleNamespace( + pp_size=pp_size, tp_group=tp_group + ) + self.assertIs( + weight_utils.get_pp_stage_load_group(), + tp_group if pp_size > 1 else _DEFAULT_LOAD_GROUP, + ) def test_iterator_rejects_unsupported_files(self): for files in (["model.pt"], ["model.bin"], ["model.safetensors", "model.pt"]): @@ -265,7 +272,6 @@ def test_iterator_rejects_unsupported_files(self): list(weight_utils.instanttensor_weights_iterator(files)) self.assertIn(files[-1], str(raised.exception)) self.safe_open.assert_not_called() - self.get_device.assert_not_called() def test_pp_embedding_can_reopen_instanttensor_checkpoint(self): from sglang.srt.speculative.pp_draft_embedding import ( @@ -277,7 +283,7 @@ def test_pp_embedding_can_reopen_instanttensor_checkpoint(self): name = "model.embed_tokens.weight" with tempfile.TemporaryDirectory() as folder: save_file({name: expected}, f"{folder}/model.safetensors") - with patch.object(loader_mod, "get_server_args", return_value=None): + with get_context().override_server_args(): hf_folder, weight_files, use_safetensors = prepare_checkpoint_files( folder, revision=None,