diff --git a/docs/docs/advanced_features/model_loading.mdx b/docs/docs/advanced_features/model_loading.mdx index eb67e3c4b09b..64696b92a26c 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. Supports CUDA and ROCm; requires PyTorch's cuda device interface. Distributed loading may retain communicator memory for the server lifetime. + pt Load weights in the PyTorch .bin format. @@ -151,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/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 9a065c0c29a7..2ebd65ea54bc 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` @@ -1574,7 +1574,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 768ceb52bad8..cb996cdceeff 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -109,6 +109,7 @@ default = true [project.optional-dependencies] checkpoint-engine = ["checkpoint-engine==0.1.2"] +instanttensor = ["instanttensor>=0.1.9"] runai = ["runai-model-streamer[s3,gcs,azure]>=0.15.7"] diffusion = [ "addict==2.4.0", @@ -170,6 +171,7 @@ test = [ "diff-cover", "expecttest", "granian>=2.6.0", + "instanttensor>=0.1.9", "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 809e33ff193b..ffb363904a04 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/python/sglang/srt/arg_groups/choices.py @@ -4,6 +4,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 f2223c1ab492..c7ba75bc36b1 100644 --- a/python/sglang/srt/arg_groups/fields/model.py +++ b/python/sglang/srt/arg_groups/fields/model.py @@ -58,6 +58,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/arg_groups/validation_hook.py b/python/sglang/srt/arg_groups/validation_hook.py index b5785fa8e6de..54ee84ae6706 100644 --- a/python/sglang/srt/arg_groups/validation_hook.py +++ b/python/sglang/srt/arg_groups/validation_hook.py @@ -239,6 +239,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/configs/load_config.py b/python/sglang/srt/configs/load_config.py index 5a0e06d026dd..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,11 +13,23 @@ 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" PT = "pt" SAFETENSORS = "safetensors" + INSTANTTENSOR = "instanttensor" NPCACHE = "npcache" DUMMY = "dummy" SHARDED_STATE = "sharded_state" @@ -50,6 +62,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 @@ -112,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 0b134976fcbc..c6b424ce8abf 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 f6cf2a775dc1..a0af2511340e 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, @@ -335,9 +339,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). @@ -1226,6 +1232,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_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py index 050af1a0befa..46b9a74afb3b 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, get_parallel from sglang.srt.utils import ( @@ -188,7 +191,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/loader.py b/python/sglang/srt/model_loader/loader.py index b0f8813b8763..a916a0c32f64 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -119,6 +119,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, @@ -359,6 +360,11 @@ def load_model( 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"} if load_format == LoadFormat.FASTSAFETENSORS: allowed_keys.add("enable_gds") @@ -490,9 +496,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"] @@ -631,6 +638,12 @@ 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, + extra_config=extra_config, + load_group=self.load_config.load_group, + ) 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 e3cca9b705ce..9fc65d822b10 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -40,7 +40,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, @@ -59,6 +63,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, @@ -1178,6 +1183,89 @@ 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, + 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": + 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( + "InstantTensor only supports .safetensors checkpoints; " + f"unsupported files: {unsupported_files}" + ) + + try: + import instanttensor + except ImportError as e: + raise ImportError( + 'Please install InstantTensor via `pip install "instanttensor>=0.1.9"`.' + ) from e + + kwargs = dict(extra_config or {}) + 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 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 + with instanttensor.safe_open( + hf_weights_files, + framework="pt", + device=device, + process_group=process_group, + copy=True, + **kwargs, + ) 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/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index c207df8f7b5a..2b441d2e4943 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 @@ -72,6 +73,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 @@ -93,6 +95,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 fda170cf1581..49ca5f002978 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -21,6 +21,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, @@ -188,6 +189,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, @@ -200,6 +202,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 debbfc7bc505..c0788cbaadb2 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -58,6 +58,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, @@ -189,6 +190,9 @@ 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_pp_stage_load_group() with ( draft_tp_context(self.draft_owns_attention), draft_pp_context(), @@ -204,6 +208,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 6f5ba5cedd5e..120098a26be8 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/python/sglang/test/runners.py b/python/sglang/test/runners.py index 757880e21be1..dd912cd4ee02 100644 --- a/python/sglang/test/runners.py +++ b/python/sglang/test/runners.py @@ -604,6 +604,8 @@ def __init__( lora_eviction_policy: str = "lru", 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" @@ -629,6 +631,8 @@ def __init__( dtype=get_dtype_str(torch_dtype), 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 new file mode 100644 index 000000000000..916253e899c6 --- /dev/null +++ b/test/registered/model_loading/test_instanttensor_loader.py @@ -0,0 +1,63 @@ +import multiprocessing as mp +import os +import unittest + +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="2-gpu-large") + +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", + tp_size=2, + disable_cuda_graph=True, + disable_radix_cache=True, + max_total_tokens=256, + ) + + 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__": + 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..240c4cac521d --- /dev/null +++ b/test/registered/unit/model_loader/test_instanttensor_loader.py @@ -0,0 +1,301 @@ +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 + +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 +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 + +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 + + 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 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.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( + 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_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, + "io_depth": 64, + "backend": ["URING", "AIO"], + } + config = LoadConfig( + 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) + resolved = loader_mod.DefaultModelLoader.ResolvedSource( + source=source, + hf_folder="model", + weight_files=("model.safetensors",), + use_safetensors=True, + ) + 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=group.device_group, + copy=True, + **{**options, "backend": [_Backend.URING, _Backend.AIO]}, + ) + self.assertEqual(config.model_loader_extra_config, options) + + 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.safe_open.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(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_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 + ), + tp_group=SimpleNamespace(world_size=1, device_group=object()), + ) + result = list( + weight_utils.instanttensor_weights_iterator( + ["model.safetensors"] + ) + ) + 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, + ) + + 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()) + 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() + list( + weight_utils.instanttensor_weights_iterator( + ["model.safetensors"], load_group=group + ) + ) + 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, + ) + + 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"]): + 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() + + 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 get_context().override_server_args(): + 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() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index b962394f4ce1..92f69b7226c2 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -186,6 +186,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",