From cf3813f4cec28c490fbdab7009031014a0fe5918 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 18:37:29 +0800 Subject: [PATCH] [diffusion] feat: add weight source reader (#35668) --- .../runtime/loader/weight_readers/__init__.py | 75 +++++++++++++++++++ .../runtime/loader/weight_readers/base.py | 47 ++++++++++++ .../loader/weight_readers/runai_streamer.py | 56 ++++++++++++++ .../loader/weight_readers/safetensors_mmap.py | 47 ++++++++++++ .../runtime/loader/weight_utils.py | 75 +++++++------------ .../test/unit/test_weight_readers.py | 68 +++++++++++++++++ 6 files changed, 321 insertions(+), 47 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/loader/weight_readers/__init__.py create mode 100644 python/sglang/multimodal_gen/runtime/loader/weight_readers/base.py create mode 100644 python/sglang/multimodal_gen/runtime/loader/weight_readers/runai_streamer.py create mode 100644 python/sglang/multimodal_gen/runtime/loader/weight_readers/safetensors_mmap.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_weight_readers.py diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_readers/__init__.py b/python/sglang/multimodal_gen/runtime/loader/weight_readers/__init__.py new file mode 100644 index 000000000..ec3028c1b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/loader/weight_readers/__init__.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Ways of reading checkpoint weights, and the rule for picking one.""" + +from sglang.multimodal_gen import envs +from sglang.multimodal_gen.runtime.loader.weight_readers.base import ( + WeightReader, +) +from sglang.multimodal_gen.runtime.loader.weight_readers.runai_streamer import ( + RunaiStreamerReader, +) +from sglang.multimodal_gen.runtime.loader.weight_readers.safetensors_mmap import ( + SafetensorsMmapReader, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# The fallback is last and always available, so selection cannot come up empty. +_READERS: tuple[type, ...] = (RunaiStreamerReader, SafetensorsMmapReader) +FALLBACK_READER = SafetensorsMmapReader + + +def available_reader_names() -> list[str]: + return [b.name for b in _READERS if b.is_available()] + + +def select_weight_reader( + *, + requested: str | None = None, + needs_key_filter: bool = False, +) -> WeightReader: + """Pick a reader, honouring an explicit request where it can be honoured. + + `requested` names a reader; None means take the environment's preference. + A reader that cannot skip keys is passed over when the caller needs to, + because reading the whole checkpoint to discard most of it is worse than + reading the part that was asked for more slowly. + """ + if requested is not None: + chosen = next((b for b in _READERS if b.name == requested), None) + if chosen is None: + raise ValueError( + f"unknown weight reader {requested!r}; " + f"available: {available_reader_names()}" + ) + elif envs.SGLANG_USE_RUNAI_MODEL_STREAMER and RunaiStreamerReader.is_available(): + chosen = RunaiStreamerReader + else: + chosen = FALLBACK_READER + + if not chosen.is_available(): + logger.info( + "Weight reader %s is not installed; using %s", + chosen.name, + FALLBACK_READER.name, + ) + chosen = FALLBACK_READER + if needs_key_filter and not chosen.supports_key_filter: + logger.debug( + "Weight reader %s cannot skip keys at load time; using %s", + chosen.name, + FALLBACK_READER.name, + ) + chosen = FALLBACK_READER + return chosen() + + +__all__ = [ + "FALLBACK_READER", + "RunaiStreamerReader", + "SafetensorsMmapReader", + "WeightReader", + "available_reader_names", + "select_weight_reader", +] diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_readers/base.py b/python/sglang/multimodal_gen/runtime/loader/weight_readers/base.py new file mode 100644 index 000000000..05e64bf31 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/loader/weight_readers/base.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +"""What a weight source has to provide, and what distinguishes one from another. + +Reading a checkpoint used to be a boolean: Run:ai streamer, or `safe_open`. The +two differ in more than speed, and the differences decide correctness and memory +behaviour rather than taste, so they are stated here as capabilities: + +``supports_key_filter`` + Whether the reader can skip keys while reading. The streamer materializes + every tensor before handing any of them over, so a caller that only wants + part of a checkpoint cannot save anything by asking it. + +``retains_file_mapping`` + Whether the tensors it yields are views into the checkpoint file. Those + pages are file-backed, so the kernel can drop them under pressure without + swap; a reader that copies into anonymous memory gives the kernel nothing + to reclaim. +""" + +from typing import Callable, ClassVar, Iterator, Protocol, runtime_checkable + +import torch + + +@runtime_checkable +class WeightReader(Protocol): + """Yields ``(name, tensor)`` for every weight in a set of checkpoint files.""" + + name: ClassVar[str] + supports_key_filter: ClassVar[bool] + retains_file_mapping: ClassVar[bool] + + @classmethod + def is_available(cls) -> bool: + """Whether this reader can run at all in this install.""" + + def iter_weights( + self, + files: list[str], + *, + device: str, + to_cpu: bool, + key_filter: Callable[[str], bool] | None = None, + clone_tensors: bool = True, + show_progress: bool = True, + ) -> Iterator[tuple[str, torch.Tensor]]: + """Iterate the weights, in whatever order the reader finds them.""" diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_readers/runai_streamer.py b/python/sglang/multimodal_gen/runtime/loader/weight_readers/runai_streamer.py new file mode 100644 index 000000000..6988a5d1f --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/loader/weight_readers/runai_streamer.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run:ai Model Streamer: fastest to read, but it copies into anonymous memory.""" + +from typing import Callable, ClassVar, Iterator + +import torch + +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +try: + from runai_model_streamer import SafetensorsStreamer + + HAS_RUNAI_MODEL_STREAMER = True +except ImportError: + SafetensorsStreamer = None + HAS_RUNAI_MODEL_STREAMER = False + + +class RunaiStreamerReader: + name: ClassVar[str] = "runai_streamer" + # it materializes every tensor before yielding, so filtering saves nothing + supports_key_filter: ClassVar[bool] = False + retains_file_mapping: ClassVar[bool] = False + + @classmethod + def is_available(cls) -> bool: + return HAS_RUNAI_MODEL_STREAMER + + def iter_weights( + self, + files: list[str], + *, + device: str, + to_cpu: bool, + key_filter: Callable[[str], bool] | None = None, + clone_tensors: bool = True, + show_progress: bool = True, + ) -> Iterator[tuple[str, torch.Tensor]]: + logger.info( + "Loading safetensors with Run:ai Model Streamer to %s", + "cpu" if to_cpu else device, + ) + with SafetensorsStreamer() as streamer: + if to_cpu: + streamer.stream_files(files) + else: + streamer.stream_files(files, device=device) + for name, tensor in streamer.get_tensors(): + if key_filter is not None and not key_filter(name): + continue + if to_cpu or clone_tensors: + yield name, tensor.clone().detach() + else: + yield name, tensor diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_readers/safetensors_mmap.py b/python/sglang/multimodal_gen/runtime/loader/weight_readers/safetensors_mmap.py new file mode 100644 index 000000000..24f3a428a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/loader/weight_readers/safetensors_mmap.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +"""`safe_open`: slower to read, and the only source whose pages stay reclaimable. + +`safe_open` maps the file, so a CPU tensor it yields is a view into the +checkpoint rather than a copy. Those pages are file-backed, which is what lets +the kernel drop them under memory pressure even on a host with no swap. +""" + +from typing import Callable, ClassVar, Iterator + +import torch +from safetensors.torch import safe_open +from tqdm.auto import tqdm + +_BAR_FORMAT = "{desc}: {percentage:.0f}%|{bar}| {n_fmt}/{total_fmt}" + + +class SafetensorsMmapReader: + name: ClassVar[str] = "safetensors" + supports_key_filter: ClassVar[bool] = True + retains_file_mapping: ClassVar[bool] = True + + @classmethod + def is_available(cls) -> bool: + return True + + def iter_weights( + self, + files: list[str], + *, + device: str, + to_cpu: bool, + key_filter: Callable[[str], bool] | None = None, + clone_tensors: bool = True, + show_progress: bool = True, + ) -> Iterator[tuple[str, torch.Tensor]]: + for path in tqdm( + files, + desc="Loading safetensors checkpoint shards", + disable=not show_progress, + bar_format=_BAR_FORMAT, + ): + with safe_open(path, framework="pt", device=device) as handle: + for name in handle.keys(): # noqa: SIM118 + if key_filter is not None and not key_filter(name): + continue + yield name, handle.get_tensor(name) diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index 103e309d2..a7aaa7d13 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py @@ -18,16 +18,16 @@ from safetensors.torch import safe_open from torch.distributed.tensor import DTensor from tqdm.auto import tqdm -try: - from runai_model_streamer import SafetensorsStreamer - - HAS_RUNAI_MODEL_STREAMER = True -except ImportError: - HAS_RUNAI_MODEL_STREAMER = False - -from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan +from sglang.multimodal_gen.runtime.loader.weight_readers import ( + FALLBACK_READER, + RunaiStreamerReader, + select_weight_reader, +) +from sglang.multimodal_gen.runtime.loader.weight_readers.runai_streamer import ( + HAS_RUNAI_MODEL_STREAMER, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -230,14 +230,19 @@ def safetensors_weights_iterator( device = str(checkpoint_device) else: device = "cpu" if to_cpu else str(get_local_torch_device()) - if use_runai_model_streamer is None: - use_runai_model_streamer = ( - HAS_RUNAI_MODEL_STREAMER and envs.SGLANG_USE_RUNAI_MODEL_STREAMER + # The caller may still pass the old boolean; map it onto a backend name so + # there is one place that decides, and it is the place that knows which + # backends can skip keys. + requested = None + if use_runai_model_streamer is not None: + requested = ( + RunaiStreamerReader.name + if use_runai_model_streamer + else FALLBACK_READER.name ) - if key_filter is not None: - # streamer filters after materializing all tensors, so it cannot skip - # a checkpoint partition at load time - use_runai_model_streamer = False + backend = select_weight_reader( + requested=requested, needs_key_filter=key_filter is not None + ) # Validate files before loading corrupted_files, duplicate_files_by_key = _scan_safetensors_files(hf_weights_files) @@ -273,38 +278,14 @@ def safetensors_weights_iterator( _raise_if_duplicate_safetensors_keys(duplicate_files_by_key) - if use_runai_model_streamer: - logger.info( - "Loading safetensors with Run:ai Model Streamer to %s", - "cpu" if to_cpu else device, - ) - with SafetensorsStreamer() as streamer: - if to_cpu: - streamer.stream_files(hf_weights_files) - else: - streamer.stream_files(hf_weights_files, device=device) - for name, tensor in streamer.get_tensors(): - if key_filter is not None and not key_filter(name): - continue - if to_cpu: - yield name, tensor.clone().detach() - elif clone_streamed_tensors: - yield name, tensor.clone().detach() - else: - yield name, tensor - else: - for st_file in tqdm( - hf_weights_files, - desc="Loading safetensors checkpoint shards", - disable=not enable_tqdm, - bar_format=_BAR_FORMAT, - ): - with safe_open(st_file, framework="pt", device=device) as f: - for name in f.keys(): # noqa: SIM118 - if key_filter is not None and not key_filter(name): - continue - param = f.get_tensor(name) - yield name, param + yield from backend.iter_weights( + hf_weights_files, + device=device, + to_cpu=to_cpu, + key_filter=key_filter, + clone_tensors=clone_streamed_tensors, + show_progress=enable_tqdm, + ) def _load_pt_file(bin_file: str, device: str) -> dict: diff --git a/python/sglang/multimodal_gen/test/unit/test_weight_readers.py b/python/sglang/multimodal_gen/test/unit/test_weight_readers.py new file mode 100644 index 000000000..4bb3a924f --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_weight_readers.py @@ -0,0 +1,68 @@ +"""Which backend reads the checkpoint, and why that choice is not a boolean.""" + +import pytest + +from sglang.multimodal_gen.runtime.loader import weight_readers +from sglang.multimodal_gen.runtime.loader.weight_readers import ( + FALLBACK_READER, + RunaiStreamerReader, + SafetensorsMmapReader, + available_reader_names, + select_weight_reader, +) + + +class TestCapabilities: + def test_the_streamer_cannot_skip_keys(self): + # it materializes every tensor before yielding any of them + assert not RunaiStreamerReader.supports_key_filter + + def test_only_the_mapping_backend_leaves_pages_reclaimable(self): + assert SafetensorsMmapReader.retains_file_mapping + assert not RunaiStreamerReader.retains_file_mapping + + def test_the_fallback_is_always_available(self): + assert FALLBACK_READER.is_available() + assert FALLBACK_READER.name in available_reader_names() + + +class TestSelection: + def test_an_explicit_request_is_honoured(self): + assert select_weight_reader(requested="safetensors").name == "safetensors" + + def test_an_unknown_name_is_an_error_not_a_silent_fallback(self): + with pytest.raises(ValueError, match="unknown weight reader"): + select_weight_reader(requested="does_not_exist") + + def test_a_key_filter_passes_over_a_backend_that_cannot_filter(self): + # reading the whole checkpoint to discard most of it is worse than + # reading the requested part more slowly + chosen = select_weight_reader(requested="runai_streamer", needs_key_filter=True) + assert chosen.name == FALLBACK_READER.name + + def test_a_key_filter_leaves_a_capable_backend_alone(self): + chosen = select_weight_reader(requested="safetensors", needs_key_filter=True) + assert chosen.name == "safetensors" + + def test_an_unavailable_backend_falls_back(self, monkeypatch): + monkeypatch.setattr( + RunaiStreamerReader, "is_available", classmethod(lambda cls: False) + ) + assert ( + select_weight_reader(requested="runai_streamer").name + == FALLBACK_READER.name + ) + + def test_the_environment_decides_when_nothing_is_requested(self, monkeypatch): + monkeypatch.setattr( + weight_readers.envs, "SGLANG_USE_RUNAI_MODEL_STREAMER", False + ) + assert select_weight_reader().name == FALLBACK_READER.name + + def test_the_environment_can_ask_for_the_streamer(self, monkeypatch): + if not RunaiStreamerReader.is_available(): + pytest.skip("run:ai model streamer is not installed") + monkeypatch.setattr( + weight_readers.envs, "SGLANG_USE_RUNAI_MODEL_STREAMER", True + ) + assert select_weight_reader().name == "runai_streamer"