diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 2941b8d92..1295e03c8 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -23,6 +23,10 @@ if TYPE_CHECKING: SGLANG_DIFFUSION_TRACE_FUNCTION: int = 0 SGLANG_DIFFUSION_DISABLE_EARLY_VAE_DECODER_CAST: bool = False SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE: bool = False + SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED: bool = False + SGLANG_DIFFUSION_DISABLE_MAPPED_DIRECT_READ: bool = False + SGLANG_DIFFUSION_DEBUG_HOST_MEMORY: bool = False + SGLANG_DIFFUSION_DEBUG_LAYERWISE_TIMING: bool = False SGLANG_DIFFUSION_DISABLE_LORA_MERGE_CACHE: bool = False SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD: str = "fork" SGLANG_DIFFUSION_TARGET_DEVICE: str = "cuda" @@ -34,6 +38,10 @@ if TYPE_CHECKING: VERBOSE: bool = False SGLANG_DIFFUSION_SERVER_DEV_MODE: bool = False SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER: bool = False + SGLANG_DIFFUSION_HOST_SPILL_DIR: str = os.path.expanduser( + "~/.cache/sglang/diffusion/host_spill" + ) + SGLANG_DIFFUSION_DISABLE_HOST_SPILL: bool = False SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB: float | None = None SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB: float | None = None SGLANG_DIFFUSION_STAGE_LOGGING: bool = False @@ -248,6 +256,16 @@ environment_variables: dict[str, Callable[[], Any]] = { "SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER": _lazy_bool( "SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER" ), + # Where transformed weight copies (fused q/k/v, reordered rows) live as + # file mappings when host copies must stay reclaimable; reused across + # starts of the same checkpoint. + "SGLANG_DIFFUSION_HOST_SPILL_DIR": _lazy_str( + "SGLANG_DIFFUSION_HOST_SPILL_DIR", + os.path.expanduser("~/.cache/sglang/diffusion/host_spill"), + ), + "SGLANG_DIFFUSION_DISABLE_HOST_SPILL": _lazy_bool( + "SGLANG_DIFFUSION_DISABLE_HOST_SPILL" + ), # Test hook: make the host memory budget behave as if the machine had this # many GiB of RAM (available = this figure minus the process's own # anonymous memory). CI uses it to exercise the constrained placement @@ -304,6 +322,28 @@ environment_variables: dict[str, Callable[[], Any]] = { "SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE": _lazy_bool( "SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE" ), + # Kill-switch: do not madvise(MADV_WILLNEED) mapped layers ahead of the + # courier; their pages arrive at fault-time readahead beats instead. + "SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED": _lazy_bool( + "SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED" + ), + # Kill-switch: on a shared host/device pool the courier reads mapped layers + # from their checkpoint files with O_DIRECT instead of through the page + # cache. This forces the mmap path. + "SGLANG_DIFFUSION_DISABLE_MAPPED_DIRECT_READ": _lazy_bool( + "SGLANG_DIFFUSION_DISABLE_MAPPED_DIRECT_READ" + ), + # Debug: after auto residency settles, log where this process's host memory + # sits -- per component and per kind (anonymous, mapped, pinned) -- next to + # the kernel's view of the process. + "SGLANG_DIFFUSION_DEBUG_HOST_MEMORY": _lazy_bool( + "SGLANG_DIFFUSION_DEBUG_HOST_MEMORY" + ), + # Debug: at the end of every layerwise stage, log where the courier and the + # compute thread spent their time (populate, memcpy, H2D, waits). + "SGLANG_DIFFUSION_DEBUG_LAYERWISE_TIMING": _lazy_bool( + "SGLANG_DIFFUSION_DEBUG_LAYERWISE_TIMING" + ), # Kill-switch: keep LoRA-merged weights in anonymous host memory instead # of the file-backed LoRA merge cache. "SGLANG_DIFFUSION_DISABLE_LORA_MERGE_CACHE": _lazy_bool( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 6d2b4c0b8..7f3460dc9 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -4,7 +4,6 @@ import os import torch import torch.nn as nn -from safetensors.torch import load_file as safetensors_load_file from safetensors.torch import safe_open from safetensors.torch import save_file as safetensors_save_file @@ -170,6 +169,22 @@ def _decode_dtype_store_path( ) +def _load_safetensors_file(path: str) -> dict: + """The VAE checkpoint itself: read-only where host copies are redundant.""" + from sglang.multimodal_gen.runtime.loader.utils import ( + _load_safetensors_file as _load, + ) + + return _load(path) + + +def _load_store(path: str) -> dict: + """Map the store read-only where host copies are redundant (see loader.utils).""" + from sglang.multimodal_gen.runtime.loader.utils import _load_safetensors_file + + return _load_safetensors_file(path) + + def _assign_matching_store(vae, mapped: dict, dtype: torch.dtype) -> bool: """Adopt a decode-dtype store if it matches the module, else refuse.""" state = vae.state_dict() @@ -199,7 +214,7 @@ def _rehome_cast_weights_to_file( path = _decode_dtype_store_path(component_model_path, component_name, dtype) try: if os.path.exists(path): - mapped = safetensors_load_file(path) + mapped = _load_store(path) if mapped and _assign_matching_store(vae, mapped, dtype): return len(mapped), True raise ValueError("existing decode-dtype store does not match the module") @@ -215,7 +230,7 @@ def _rehome_cast_weights_to_file( tmp = f"{path}.tmp.{os.getpid()}" safetensors_save_file({k: v.contiguous() for k, v in cast_state.items()}, tmp) os.replace(tmp, path) - mapped = safetensors_load_file(path) + mapped = _load_store(path) if set(mapped) != set(cast_state): raise ValueError("decode-dtype store does not match the cast weights") vae.load_state_dict(mapped, strict=False, assign=True) @@ -671,7 +686,7 @@ class VAELoader(WeightOverrideComponentLoader): loaded = {} for sf_path in safetensors_list: - loaded.update(safetensors_load_file(sf_path)) + loaded.update(_load_safetensors_file(sf_path)) _backfill_ltx2_audio_vae_latent_stats(loaded, component_type) num_deparameterized = adopt_plain_weight_norm_state(vae, loaded) target_state = vae.state_dict() diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index c8b01d1d6..2f7b737b7 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -38,6 +38,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import ( split_bitsandbytes_4bit_state, ) from sglang.multimodal_gen.runtime.loader import rank_local_checkpoint +from sglang.multimodal_gen.runtime.loader.host_spill import HostSpill from sglang.multimodal_gen.runtime.loader.utils import ( finalize_loaded_model, get_param_names_mapping, @@ -48,6 +49,9 @@ from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan from sglang.multimodal_gen.runtime.loader.weight_utils import ( safetensors_weights_iterator, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + host_copies_are_redundant, +) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.quantization_utils import ( @@ -315,6 +319,18 @@ def maybe_load_fsdp_model( # layerwise offload replaces block parameters with placeholders after # load, so compatible checkpoint tensors stay file-backed on CPU model._keep_checkpoint_mapping = True + host_spill = None + if ( + weight_dir_list + and weights_iterator is None + and weight_load_plan.checkpoint_load_device.type == "cpu" + and host_copies_are_redundant() + ): + # Compatible tensors stay on the checkpoint mapping; the fused and + # reordered ones are materialized. On a shared pool their anonymous + # copies are never reclaimable (10.4 GiB for the H3 DiT's fused + # q/k/v), so they are materialized into file mappings instead. + host_spill = HostSpill.for_checkpoint(weight_dir_list) defer_cpu_placement = bool( component_starts_on_cpu and weight_load_plan.defer_cpu_placement @@ -437,7 +453,10 @@ def maybe_load_fsdp_model( weight_load_plan.load_full_state_dict_on_device ), preconverted_state_dict=preconverted_state_dict, + host_spill=host_spill, ) + if host_spill is not None: + host_spill.log_summary(type(model).__name__) if bnb_quant_states: attach_bitsandbytes_4bit_quant_states( dict(model.named_parameters()), bnb_quant_states @@ -557,6 +576,7 @@ def load_model_from_full_model_state_dict( | None ) = None, allow_device_tensor_assignment: bool = False, + host_spill: HostSpill | None = None, ) -> _IncompatibleKeys: """ Converting full state dict into a sharded state dict @@ -589,7 +609,12 @@ def load_model_from_full_model_state_dict( full_sd_iterator, param_names_mapping, valid_target_names=set(meta_sd.keys()), + fused_tensor_factory=(None if host_spill is None else host_spill.tensor), ) # type: ignore + if host_spill is not None: + for name, tensor in custom_param_sd.items(): + if isinstance(tensor, torch.Tensor): + host_spill.seal(name, tensor.shape, tensor.dtype) else: custom_param_sd, reverse_param_names_mapping = preconverted_state_dict @@ -734,11 +759,21 @@ def load_model_from_full_model_state_dict( ): sharded_tensor = full_tensor else: - sharded_tensor = torch.empty_like( - meta_sharded_param, - device=checkpoint_load_device, - dtype=target_dtype, + spilled = ( + None + if host_spill is None or checkpoint_load_device.type != "cpu" + else host_spill.tensor( + target_param_name, meta_sharded_param.shape, target_dtype + ) ) + if spilled is not None: + sharded_tensor = spilled[0] + else: + sharded_tensor = torch.empty_like( + meta_sharded_param, + device=checkpoint_load_device, + dtype=target_dtype, + ) # Preserve requires_grad flag to avoid errors with non-floating dtypes requires_grad = meta_sharded_param.requires_grad temp_param = _make_param_like(actual_param, sharded_tensor) @@ -759,6 +794,10 @@ def load_model_from_full_model_state_dict( f"param_cls={type(actual_param).__name__}" ) from exc sharded_tensor = temp_param.data + if host_spill is not None and spilled is not None: + host_spill.seal( + target_param_name, meta_sharded_param.shape, target_dtype + ) else: # In cases where parts of the model aren't sharded, some parameters will be plain tensors sharded_tensor = full_tensor diff --git a/python/sglang/multimodal_gen/runtime/loader/host_spill.py b/python/sglang/multimodal_gen/runtime/loader/host_spill.py new file mode 100644 index 000000000..f5d28b49a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/loader/host_spill.py @@ -0,0 +1,158 @@ +"""File-backed homes for the host copies a checkpoint mapping cannot provide. + +A weight that is fused (q/k/v into one projection), sharded or otherwise +transformed at load has no checkpoint bytes to stay mapped on, so the loader +materializes it. Anonymous memory is the wrong home for that copy on a host +that keeps everything else mapped: it is never reclaimable, and on a shared +CPU/GPU pool it is memory the page cache and the device both lose. A shared +file mapping under the cache directory holds the same bytes as page cache +instead -- reclaimable, readable with O_DIRECT, and, once written, reusable +by the next start of the same checkpoint. +""" + +from __future__ import annotations + +import hashlib +import os +import shutil +from pathlib import Path +from typing import Callable, Iterable + +import torch + +from sglang.multimodal_gen import envs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Below this a copy is not worth a file: the bookkeeping costs more than the +# bytes it would return to the pool. +MIN_SPILL_BYTES = 64 << 20 +# Free space to leave on the spill filesystem after a write. +SPILL_DISK_RESERVE_BYTES = 2 << 30 + +FusedTensorFactory = Callable[ + [str, torch.Size, torch.dtype], "tuple[torch.Tensor, bool] | None" +] + + +def checkpoint_fingerprint(weight_dirs: Iterable[str]) -> str: + """Identity of a checkpoint on disk: each shard's path, size and mtime.""" + digest = hashlib.sha1() + for weight_dir in sorted(str(d) for d in weight_dirs): + root = Path(weight_dir) + files = [root] if root.is_file() else sorted(root.glob("*.safetensors")) + for path in files: + try: + stat = path.stat() + except OSError: + continue + digest.update( + f"{path.resolve()}|{stat.st_size}|{stat.st_mtime_ns}\n".encode() + ) + return digest.hexdigest()[:20] + + +class HostSpill: + """Hands out file-backed tensors keyed by (parameter, dtype, shape). + + A tensor comes back ``(tensor, filled)``: ``filled`` says an earlier run + wrote and sealed the same key, so the caller can skip producing it. A + caller that produces the bytes must ``seal`` the key afterwards; an + unsealed file is treated as garbage and rewritten. + """ + + def __init__(self, directory: str | os.PathLike[str], fingerprint: str): + self.directory = Path(directory) / fingerprint + self._disabled_reason: str | None = None + self.bytes_written = 0 + self.bytes_reused = 0 + self.count_written = 0 + self.count_reused = 0 + self._open: dict[str, str] = {} + + @classmethod + def for_checkpoint(cls, weight_dirs: Iterable[str]) -> HostSpill | None: + if envs.SGLANG_DIFFUSION_DISABLE_HOST_SPILL: + return None + directory = os.path.expanduser(envs.SGLANG_DIFFUSION_HOST_SPILL_DIR) + return cls(directory, checkpoint_fingerprint(weight_dirs)) + + def _path(self, key: str) -> Path: + return self.directory / (hashlib.sha1(key.encode()).hexdigest() + ".bin") + + def _disable(self, reason: str) -> None: + if self._disabled_reason is None: + self._disabled_reason = reason + logger.warning( + "Host spill disabled for this load: %s; transformed weights " + "fall back to anonymous memory.", + reason, + ) + + def tensor( + self, name: str, shape: torch.Size, dtype: torch.dtype + ) -> tuple[torch.Tensor, bool] | None: + """A file-backed tensor for ``name``, or None to use anonymous memory.""" + if self._disabled_reason is not None: + return None + numel = 1 + for dim in shape: + numel *= int(dim) + nbytes = numel * torch.empty((), dtype=dtype).element_size() + if nbytes < MIN_SPILL_BYTES: + return None + key = f"{name}|{dtype}|{tuple(int(d) for d in shape)}" + path = self._path(key) + sealed = path.with_suffix(".ok") + try: + self.directory.mkdir(parents=True, exist_ok=True) + filled = sealed.exists() and path.exists() and path.stat().st_size == nbytes + if not filled: + sealed.unlink(missing_ok=True) + free = shutil.disk_usage(self.directory).free + if free < nbytes + SPILL_DISK_RESERVE_BYTES: + self._disable( + f"{free / 2**30:.1f} GiB free under {self.directory}, " + f"{nbytes / 2**30:.1f} GiB needed" + ) + return None + storage = torch.from_file(str(path), shared=True, size=numel, dtype=dtype) + except (OSError, RuntimeError) as exc: + self._disable(f"{type(exc).__name__}: {exc}") + return None + tensor = storage.view(tuple(int(d) for d in shape)) + if filled: + self.bytes_reused += nbytes + self.count_reused += 1 + else: + self._open[key] = str(sealed) + self.bytes_written += nbytes + self.count_written += 1 + return tensor, filled + + def seal(self, name: str, shape: torch.Size, dtype: torch.dtype) -> None: + """Mark a key as completely written so the next start can reuse it.""" + key = f"{name}|{dtype}|{tuple(int(d) for d in shape)}" + sealed = self._open.pop(key, None) + if sealed is None: + return + try: + with open(sealed, "w") as handle: + handle.write("ok\n") + except OSError as exc: + logger.debug("could not seal %s: %s", sealed, exc) + + def log_summary(self, component: str) -> None: + if self.count_written == 0 and self.count_reused == 0: + return + logger.info( + "%s: %d transformed weights (%.2f GiB) live in file mappings under %s " + "(%d written, %d reused from an earlier start).", + component, + self.count_written + self.count_reused, + (self.bytes_written + self.bytes_reused) / 2**30, + self.directory, + self.count_written, + self.count_reused, + ) diff --git a/python/sglang/multimodal_gen/runtime/loader/readonly_safetensors.py b/python/sglang/multimodal_gen/runtime/loader/readonly_safetensors.py new file mode 100644 index 000000000..c2284d1f3 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/loader/readonly_safetensors.py @@ -0,0 +1,95 @@ +"""Read-only mappings of safetensors files. + +safetensors maps a file for torch through ``UntypedStorage.from_file(shared=False)``: +a private, writable mapping. On a shared CPU/GPU pool that permission costs +memory: when the device copies from such a mapping the driver pins the pages +with write intent, the kernel breaks copy-on-write, and every page copied +becomes anonymous memory -- 1 GiB copied in 1 GiB of unreclaimable RAM, at a +fraction of the bandwidth (0.1-3.8 GiB/s against 27 GiB/s on a GB10). A +read-only mapping copies at full speed and leaves the pages as page cache. + +Frozen weights never need the write permission; a write into one here faults +instead of silently copying the page, which is the invariant we want. +""" + +from __future__ import annotations + +import json +import mmap +import os +import struct +import warnings +from typing import Iterator + +import torch + +_DTYPES = { + "F64": torch.float64, + "F32": torch.float32, + "F16": torch.float16, + "BF16": torch.bfloat16, + "I64": torch.int64, + "I32": torch.int32, + "I16": torch.int16, + "I8": torch.int8, + "U8": torch.uint8, + "BOOL": torch.bool, + "F8_E4M3": torch.float8_e4m3fn, + "F8_E5M2": torch.float8_e5m2, +} + +# Mappings stay for the life of the process: the tensors handed out are views +# into them, and the layerwise manager keeps such views as its host store. +_MAPPINGS: dict[str, mmap.mmap] = {} + + +def _mapping(path: str) -> mmap.mmap: + real = os.path.realpath(path) + mapped = _MAPPINGS.get(real) + if mapped is None: + fd = os.open(real, os.O_RDONLY) + try: + size = os.fstat(fd).st_size + mapped = mmap.mmap(fd, size, prot=mmap.PROT_READ, flags=mmap.MAP_PRIVATE) + finally: + os.close(fd) + _MAPPINGS[real] = mapped + return mapped + + +def _header(mapped: mmap.mmap) -> tuple[dict, int]: + (n,) = struct.unpack(" torch.Tensor: + dtype = _DTYPES[meta["dtype"]] + start, end = meta["data_offsets"] + shape = tuple(meta["shape"]) + if end == start: + return torch.empty(shape, dtype=dtype) + count = (end - start) // torch.empty((), dtype=dtype).element_size() + with warnings.catch_warnings(): + # torch warns that the buffer is not writable; that is the point. + warnings.simplefilter("ignore") + flat = torch.frombuffer(mapped, dtype=dtype, count=count, offset=base + start) + return flat.view(shape) + + +def safetensors_keys(path: str) -> list[str]: + header, _ = _header(_mapping(path)) + return [name for name in header if name != "__metadata__"] + + +def iter_safetensors_readonly(path: str) -> Iterator[tuple[str, torch.Tensor]]: + """(name, tensor) for every tensor in the file, as views of a read-only mapping.""" + mapped = _mapping(path) + header, base = _header(mapped) + for name, meta in header.items(): + if name == "__metadata__": + continue + yield name, _tensor(mapped, base, meta) + + +def load_safetensors_readonly(path: str) -> dict[str, torch.Tensor]: + return dict(iter_safetensors_readonly(path)) diff --git a/python/sglang/multimodal_gen/runtime/loader/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index 6a601d6d0..15036492f 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -181,10 +181,45 @@ def get_param_names_mapping( return mapping_fn +def _fuse_tensors( + target_param_name: str, + tensors: list[torch.Tensor], + fused_tensor_factory: ( + Callable[[str, torch.Size, torch.dtype], tuple[torch.Tensor, bool] | None] + | None + ), +) -> torch.Tensor: + """Concatenate the pieces of one parameter along dim 0. + + The factory, when given, provides the destination (a file mapping that + outlives anonymous memory) and says whether an earlier run already filled + it -- then the pieces are not even read. + """ + if fused_tensor_factory is None or any(t.device.type != "cpu" for t in tensors): + return torch.cat(tensors, dim=0) + if ( + len({tuple(t.shape[1:]) for t in tensors}) != 1 + or len({t.dtype for t in tensors}) != 1 + ): + return torch.cat(tensors, dim=0) + shape = torch.Size([sum(t.shape[0] for t in tensors), *tensors[0].shape[1:]]) + provided = fused_tensor_factory(target_param_name, shape, tensors[0].dtype) + if provided is None: + return torch.cat(tensors, dim=0) + out, filled = provided + if not filled: + torch.cat(tensors, dim=0, out=out) + return out + + def hf_to_custom_state_dict( hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]], param_names_mapping: Callable[[str], tuple[str, Any, Any]], valid_target_names: set[str] | None = None, + fused_tensor_factory: ( + Callable[[str, torch.Size, torch.dtype], tuple[torch.Tensor, bool] | None] + | None + ) = None, *, strict: bool = False, ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]: @@ -236,7 +271,9 @@ def hf_to_custom_state_dict( to_merge_params[target_param_name][i] for i in range(num_params_to_merge) ] - full_tensor = torch.cat(sorted_tensors, dim=0) + full_tensor = _fuse_tensors( + target_param_name, sorted_tensors, fused_tensor_factory + ) del to_merge_params[target_param_name] else: continue @@ -365,10 +402,18 @@ def keep_checkpoint_mapped(*, weight_bytes: int, component: str) -> bool: choice -- its pages are resident, where a mapping's first use pays a fault. """ from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + host_copies_are_redundant, host_copies_would_not_fit, host_memory_available_bytes, ) + if host_copies_are_redundant(): + logger.info( + "%s stays on its checkpoint mapping: host and device share one " + "memory pool, so a copy would hold the same bytes twice.", + component, + ) + return True if not host_copies_would_not_fit(weight_bytes): return False logger.info( @@ -449,6 +494,27 @@ def _list_safetensors_files( return filter_duplicate_precision_variant_safetensors(found) +def _load_safetensors_file(path: str) -> dict[str, torch.Tensor]: + """One safetensors file; a read-only mapping where host copies are redundant. + + safetensors maps for torch through a private *writable* mapping, and on a + shared CPU/GPU pool a device copy from such a mapping copies every page it + touches into anonymous memory (the driver pins with write intent). A + read-only mapping copies at full speed and stays page cache. + """ + from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + host_copies_are_redundant, + ) + + if host_copies_are_redundant(): + from sglang.multimodal_gen.runtime.loader.readonly_safetensors import ( + load_safetensors_readonly, + ) + + return load_safetensors_readonly(path) + return safetensors_load_file(path) + + def load_safetensors_state_dict(model_path: str) -> dict[str, torch.Tensor]: """Load one safetensors checkpoint, including an indexed sharded set.""" index_path = _select_safetensors_index_file(model_path, _DEFAULT_SAFETENSORS_INDEX) @@ -456,7 +522,7 @@ def load_safetensors_state_dict(model_path: str) -> dict[str, torch.Tensor]: if index_path is not None: state_dict: dict[str, torch.Tensor] = {} for path in safetensors_files: - state_dict.update(safetensors_load_file(path)) + state_dict.update(_load_safetensors_file(path)) return state_dict if not safetensors_files: @@ -466,7 +532,7 @@ def load_safetensors_state_dict(model_path: str) -> dict[str, torch.Tensor]: f"Found {len(safetensors_files)} safetensors files in {model_path} " "and no index to disambiguate them." ) - return safetensors_load_file(safetensors_files[0]) + return _load_safetensors_file(safetensors_files[0]) BYTES_PER_GB = 1024**3 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 index 24f3a428a..340427c86 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_readers/safetensors_mmap.py +++ b/python/sglang/multimodal_gen/runtime/loader/weight_readers/safetensors_mmap.py @@ -4,6 +4,11 @@ `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. + +Where host copies are redundant (the device shares the host pool) the mapping +is made read-only instead: safetensors maps writable, and a device copy from a +writable private mapping there turns every page it touches into anonymous +memory at a fraction of the bandwidth (see readonly_safetensors). """ from typing import Callable, ClassVar, Iterator @@ -12,6 +17,13 @@ import torch from safetensors.torch import safe_open from tqdm.auto import tqdm +from sglang.multimodal_gen.runtime.loader.readonly_safetensors import ( + iter_safetensors_readonly, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + host_copies_are_redundant, +) + _BAR_FORMAT = "{desc}: {percentage:.0f}%|{bar}| {n_fmt}/{total_fmt}" @@ -40,6 +52,12 @@ class SafetensorsMmapReader: disable=not show_progress, bar_format=_BAR_FORMAT, ): + if device == "cpu" and host_copies_are_redundant(): + for name, tensor in iter_safetensors_readonly(path): + if key_filter is not None and not key_filter(name): + continue + yield name, tensor + continue 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): diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 7984fd51e..a616e7230 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -6,6 +6,7 @@ from typing import Mapping, MutableMapping, Protocol, Sequence import torch import torch.nn as nn +from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD, @@ -650,6 +651,23 @@ class ComponentResidencyManager: ) self._completed_warmup_phase_peaks = dict(self._warmup_phase_peaks) self._track_warmup_memory = False + if ( + current_platform.device_shares_host_memory() + and torch.get_device_module().is_available() + ): + # One pool: every byte the caching allocator keeps reserved between + # requests is page cache the next request's streamed encoder cannot use. + torch.get_device_module().empty_cache() + if envs.SGLANG_DIFFUSION_DEBUG_HOST_MEMORY: + from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_breakdown import ( + log_host_memory_breakdown, + ) + + self._debug_requests_seen = getattr(self, "_debug_requests_seen", 0) + 1 + log_host_memory_breakdown( + self.placement_modules(), + label=f"after request {self._debug_requests_seen}", + ) def _begin_warmup_phase( self, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py index ee60f09e1..326a14a90 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py @@ -9,10 +9,20 @@ import torch.nn as nn from torch.distributed.fsdp import FSDPModule from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + shared_pool_available_bytes, +) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +# Device growth between a component's stages on a shared pool: activations and +# the allocator's reserve (9.4 GiB measured for H3 at 1344x768x124f) plus margin. +SHARED_POOL_NEXT_STAGE_HEADROOM_BYTES = 12 * 1024**3 + +logger = init_logger(__name__) if TYPE_CHECKING: from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( @@ -174,7 +184,15 @@ class ComponentOffloadStrategy(ComponentResidencyStrategy): self.wait_for_use(module, use, state) tensor = _module_reference_tensor(module) if tensor is not None and tensor.device.type != "cpu": - module.to("cpu", non_blocking=True) + # A non-blocking device->host move lands in pinned host memory the + # size of the component. On a shared pool that pins a second copy + # of the weights next to the device copy still being read from + # -- a 57 GiB DiT took 43 GiB of shared memory in under a minute + # and exhausted a GB10. Take the synchronous, pageable path there. + module.to( + "cpu", + non_blocking=not current_platform.device_shares_host_memory(), + ) self._ready_events.pop(use.component_name, None) def finish_request( @@ -231,6 +249,35 @@ class LayerwiseOffloadStrategy(ComponentResidencyStrategy): torch.mps.synchronize() module.restore_mps_cpu_non_layer_weights() torch.mps.empty_cache() + elif ( + current_platform.is_cuda() and current_platform.device_shares_host_memory() + ): + # The stage's streamed layer windows are freed but still reserved + # by the caching allocator. On a shared pool that reserve is host + # memory the next stage's mapping needs as page cache; hand it back. + empty_cache = getattr(torch.get_device_module(), "empty_cache", None) + if empty_cache is not None: + empty_cache() + # And this component's own pages are now the least valuable in the + # cache until its next stage; say so before the next phase evicts. + # The room the cache will have for this component's next stream is + # what is available now less what the stages in between need. + room_bytes = max( + 0, shared_pool_available_bytes() - SHARED_POOL_NEXT_STAGE_HEADROOM_BYTES + ) + paged_out = 0 + for manager in module.layerwise_offload_managers: + advise_cold = getattr(manager, "advise_mapped_pages_cold", None) + if advise_cold is not None: + paged_out += int(advise_cold(room_bytes=room_bytes) or 0) + if paged_out: + logger.info( + "Layerwise offload: paged out the first %.1f GiB of %s so the " + "next request's stream fits the %.1f GiB the cache can give it.", + paged_out / 1024**3, + use.component_name, + room_bytes / 1024**3, + ) def finish_request( self, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_breakdown.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_breakdown.py new file mode 100644 index 000000000..8d91989df --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_breakdown.py @@ -0,0 +1,348 @@ +"""Where a worker's host memory sits, by component and by kind. + +Debug aid behind ``SGLANG_DIFFUSION_DEBUG_HOST_MEMORY``: on a shared +host/device pool every anonymous byte the runtime keeps is a byte the page +cache cannot hold, so the breakdown says what to cut. +""" + +from __future__ import annotations + +import gc +import logging +from bisect import bisect_right +from collections.abc import Mapping + +import torch + +logger = logging.getLogger(__name__) + +GIB = 1024**3 + + +def _file_backed_ranges() -> list[tuple[int, int]]: + ranges: list[tuple[int, int]] = [] + try: + with open("/proc/self/maps") as handle: + for line in handle: + fields = line.split() + if len(fields) < 6 or fields[5].startswith("["): + continue + start, end = fields[0].split("-") + ranges.append((int(start, 16), int(end, 16))) + except OSError: + return [] + ranges.sort() + return ranges + + +def _kind( + tensor: torch.Tensor, starts: list[int], ranges: list[tuple[int, int]] +) -> str: + if tensor.is_pinned(): + return "pinned" + ptr = tensor.data_ptr() + index = bisect_right(starts, ptr) - 1 + if index >= 0 and ranges[index][0] <= ptr < ranges[index][1]: + return "mapped" + return "anonymous" + + +def _anon_vmas(min_bytes: int = 128 * 1024**2) -> list[tuple[int, int, int, str]]: + """(start, end, anonymous_bytes, vmflags) of anonymous mappings holding at least min_bytes.""" + out: list[tuple[int, int, int, str]] = [] + try: + start = end = 0 + path = "" + anon = 0 + flags = "" + with open("/proc/self/smaps") as handle: + for line in handle: + if line[0] in "0123456789abcdef" and "-" in line.split()[0]: + if path in ("", "[anon]") and anon >= min_bytes: + out.append((start, end, anon, flags)) + fields = line.split() + start, end = (int(x, 16) for x in fields[0].split("-")) + path = fields[5] if len(fields) >= 6 else "" + anon = 0 + flags = "" + elif line.startswith("Anonymous:"): + anon = int(line.split()[1]) * 1024 + elif line.startswith("VmFlags:"): + flags = line.split(":", 1)[1].strip() + if path in ("", "[anon]") and anon >= min_bytes: + out.append((start, end, anon, flags)) + except OSError: + pass + return sorted(out, key=lambda item: -item[2]) + + +def _file_mapping_cow(top: int = 3) -> tuple[int, list[tuple[str, int]]]: + """Anonymous bytes inside file mappings: pages a write copied out of the file. + + A private writable mapping (safetensors' from_file) stays "file-backed" by + pointer, but every page written to it becomes anonymous memory that no + longer drops under pressure. Returns (total, [(path, bytes)] for the largest). + """ + per_path: dict[str, int] = {} + try: + path = "" + with open("/proc/self/smaps") as handle: + for line in handle: + if line[0] in "0123456789abcdef" and "-" in line.split()[0]: + fields = line.split() + path = fields[5] if len(fields) >= 6 else "" + if path.startswith("[") or path.startswith("/dev/"): + path = "" + elif path and line.startswith("Anonymous:"): + anon = int(line.split()[1]) * 1024 + if anon: + per_path[path] = per_path.get(path, 0) + anon + except OSError: + return 0, [] + ranked = sorted(per_path.items(), key=lambda item: -item[1]) + return sum(per_path.values()), ranked[:top] + + +def _mallinfo() -> dict[str, float]: + try: + import ctypes + import ctypes.util + + libc = ctypes.CDLL(ctypes.util.find_library("c")) + + class MallInfo2(ctypes.Structure): + _fields_ = [ + (name, ctypes.c_size_t) + for name in ( + "arena", + "ordblks", + "smblks", + "hblks", + "hblkhd", + "usmblks", + "fsmblks", + "uordblks", + "fordblks", + "keepcost", + ) + ] + + libc.mallinfo2.restype = MallInfo2 + info = libc.mallinfo2() + return { + "glibc_arena": info.arena / GIB, + "glibc_mmapped": info.hblkhd / GIB, + "glibc_in_use": info.uordblks / GIB, + "glibc_free": info.fordblks / GIB, + } + except Exception: + return {} + + +def _sample_anon_vma(start: int, end: int) -> str: + """A few bytes from the start and the middle of a mapping, as hex, plus a + guess at what they hold (zeros, bf16-looking, other).""" + import ctypes + + out = [] + for offset in (0, (end - start) // 2 & ~4095): + try: + raw = ctypes.string_at(start + offset, 32) + except Exception: + out.append("unreadable") + continue + if not any(raw): + kind = "zeros" + else: + # bf16 weights: high bytes cluster around 0x3c-0x40 / 0xbc-0xc0 + highs = raw[1::2] + kind = ( + "bf16-like" + if sum(1 for b in highs if 0x38 <= (b & 0x7F) <= 0x42) >= 10 + else "other" + ) + out.append(f"@{offset:#x}:{raw[:16].hex()}({kind})") + return " ".join(out) + + +def _tensor_ptrs_inside(start: int, end: int) -> tuple[int, float]: + count = 0 + total = 0.0 + for obj in gc.get_objects(): + if isinstance(obj, torch.Tensor) and obj.device.type == "cpu": + try: + ptr = obj.data_ptr() + nbytes = obj.numel() * obj.element_size() + except Exception: + continue + if start <= ptr < end: + count += 1 + total += nbytes / GIB + return count, total + + +def _smaps_rollup() -> dict[str, float]: + totals: dict[str, float] = {} + try: + with open("/proc/self/smaps_rollup") as handle: + for line in handle: + key, _, rest = line.partition(":") + if key in ( + "Rss", + "Anonymous", + "Rss_File", + "Rss_Shmem", + "Private_Dirty", + ): + totals[key] = int(rest.split()[0]) * 1024 / GIB + except OSError: + pass + return totals + + +def host_memory_breakdown(modules: Mapping[str, object]) -> dict[str, dict[str, float]]: + """GiB of CPU tensor storage per component and kind; ``other`` covers + tensors no module owns (staging buffers, caches, activations kept alive).""" + ranges = _file_backed_ranges() + starts = [start for start, _ in ranges] + owners: dict[int, str] = {} + for name, module in modules.items(): + if not isinstance(module, torch.nn.Module): + continue + for tensor in list(module.parameters()) + list(module.buffers()): + if tensor.device.type == "cpu": + owners[tensor.untyped_storage().data_ptr()] = name + seen: set[int] = set() + table: dict[str, dict[str, float]] = {} + for obj in gc.get_objects(): + if not isinstance(obj, torch.Tensor) or obj.device.type != "cpu": + continue + try: + storage = obj.untyped_storage() + key = storage.data_ptr() + nbytes = storage.nbytes() + except Exception: + continue + if key in seen or nbytes == 0: + continue + seen.add(key) + owner = owners.get(key, "other") + kind = _kind(obj, starts, ranges) + table.setdefault(owner, {}) + table[owner][kind] = table[owner].get(kind, 0.0) + nbytes / GIB + return table + + +def log_host_memory_breakdown(modules: Mapping[str, object], *, label: str) -> None: + table = host_memory_breakdown(modules) + rollup = _smaps_rollup() + lines = [f"Host memory breakdown ({label}):"] + if rollup: + lines.append( + " process: " + + " ".join(f"{key}={value:.2f}GiB" for key, value in sorted(rollup.items())) + ) + for owner in sorted(table, key=lambda name: -sum(table[name].values())): + kinds = " ".join( + f"{kind}={value:.2f}GiB" for kind, value in sorted(table[owner].items()) + ) + lines.append(f" {owner}: {kinds}") + cow_total, cow_top = _file_mapping_cow() + if cow_total: + lines.append( + f" copy-on-write pages in file mappings: {cow_total / GIB:.2f}GiB; " + + ", ".join( + f"{'/'.join(path.rsplit('/', 3)[-3:])}={size / GIB:.2f}GiB" + for path, size in cow_top + ) + ) + device = torch.get_device_module() + if hasattr(device, "memory_allocated"): + lines.append( + f" device: allocated={device.memory_allocated() / GIB:.2f}GiB " + f"reserved={device.memory_reserved() / GIB:.2f}GiB" + ) + snapshot = getattr(device, "memory_snapshot", None) + if callable(snapshot): + try: + segments = snapshot() + except Exception: + segments = [] + ranges = [ + ( + int(s.get("address", 0)), + int(s.get("address", 0)) + int(s.get("total_size", 0)), + ) + for s in segments + ] + malloc = _mallinfo() + if malloc: + lines.append( + " glibc: " + " ".join(f"{k}={v:.2f}GiB" for k, v in malloc.items()) + ) + host_stats = getattr(device, "host_memory_stats", None) + if callable(host_stats): + try: + hs = host_stats() + big = { + k: v + for k, v in hs.items() + if isinstance(v, (int, float)) and v >= 64 * 1024**2 + } + lines.append( + " pinned host allocator: " + + ( + " ".join( + f"{k}={v / GIB:.2f}GiB" for k, v in sorted(big.items()) + ) + or "no counter >= 64 MiB" + ) + ) + except Exception as exc: + lines.append(f" pinned host allocator: unavailable ({exc})") + for start, end, anon, flags in _anon_vmas()[:16]: + inside = any(a <= start < b for a, b in ranges) + count, held = _tensor_ptrs_inside(start, end) + lines.append( + f" anon vma {start:#x}-{end:#x}: {anon / GIB:.2f}GiB " + f"{'INSIDE cuda segment' if inside else 'outside cuda segments'} [{flags}] " + f"live cpu tensors inside={count} ({held:.2f}GiB) sample: {_sample_anon_vma(start, end)}" + ) + try: + config = torch.__config__.show() + lines.append( + " torch build: " + + ", ".join( + line.strip() + for line in config.splitlines() + if "MIMALLOC" in line or "ALLOC" in line.upper() and "USE_" in line + )[:300] + ) + except Exception: + pass + big = sorted( + ( + (int(segment.get("address", 0)), int(segment.get("total_size", 0))) + for segment in segments + if int(segment.get("total_size", 0)) >= 256 * 1024**2 + ), + key=lambda item: -item[1], + )[:12] + lines.append( + " device segments >= 256 MiB: " + + ", ".join(f"{address:#x}:{size / GIB:.2f}GiB" for address, size in big) + ) + logger.info("\n".join(lines)) + + +def log_anon_vmas(label: str) -> None: + """Debug: the process's large anonymous mappings right now, for a timeline of + where they appear during startup.""" + rollup = _smaps_rollup() + lines = [ + f"Anonymous memory timeline ({label}): " + + " ".join(f"{k}={v:.2f}GiB" for k, v in sorted(rollup.items())) + ] + for start, end, anon, flags in _anon_vmas()[:8]: + lines.append(f" anon vma {start:#x}-{end:#x}: {anon / GIB:.2f}GiB [{flags}]") + logger.info("\n".join(lines)) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py index 3c6293f1e..e5d9a1f88 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py @@ -18,6 +18,7 @@ import os import psutil from sglang.multimodal_gen import envs +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -103,13 +104,36 @@ def _cgroup_dirs(mount: str) -> list[str]: return dirs -def cgroup_memory_limit_bytes() -> tuple[int, int] | None: +# memory.stat keys for the page cache a cgroup is charged for: v2, then v1. +_CGROUP_FILE_CACHE_KEYS = ("file", "cache") + + +def _cgroup_file_cache_bytes(directory: str) -> int: + try: + with open(os.path.join(directory, "memory.stat")) as handle: + for line in handle: + key, _, value = line.partition(" ") + if key in _CGROUP_FILE_CACHE_KEYS: + return int(value) + except (OSError, ValueError): + pass + return 0 + + +def cgroup_memory_limit_bytes( + *, exclude_file_cache: bool = False +) -> tuple[int, int] | None: """This process's (cap, usage) under its cgroup, or None when uncapped. The tightest cap in the chain wins. A nested cgroup -- a systemd scope with MemoryMax, a container started with --cgroup-parent -- holds this process below whatever the mount root allows, and planning against the root would commit memory the process cannot have. + + A cgroup is charged for the page cache it touches, so its usage grows by + the whole checkpoint the process maps. ``exclude_file_cache`` reports the + anonymous share alone, for callers that may spend cache the kernel would + reclaim under the cap anyway. """ for mount, limit_name, usage_name in _CGROUP_MOUNTS: tightest = None @@ -119,7 +143,10 @@ def cgroup_memory_limit_bytes() -> tuple[int, int] | None: continue if tightest is not None and limit >= tightest[0]: continue - tightest = (limit, _read_int(os.path.join(directory, usage_name)) or 0) + usage = _read_int(os.path.join(directory, usage_name)) or 0 + if exclude_file_cache: + usage = max(0, usage - _cgroup_file_cache_bytes(directory)) + tightest = (limit, usage) if tightest is not None: return tightest return None @@ -155,6 +182,33 @@ def host_memory_available_bytes() -> int: return min(available, max(0, limit - usage)) +def shared_pool_available_bytes() -> int: + """Bytes a shared host/device pool can still give this process. + + The device's own free figure is the kernel's MemFree, which leaves out the + page cache -- memory the kernel hands back on demand and a placement may + therefore spend. A cgroup cap is honoured on its anonymous share only, for + the same reason: the cache charged to the cgroup is reclaimed under the cap. + """ + available = int(psutil.virtual_memory().available) + capped = cgroup_memory_limit_bytes(exclude_file_cache=True) + if capped is None: + return available + limit, anonymous = capped + return min(available, max(0, limit - anonymous)) + + +def host_copies_are_redundant() -> bool: + """Whether a host copy of a mapped weight buys nothing. + + When host and device share one physical pool the device reads page-cache + pages directly, so a pinned or pageable copy holds the same bytes twice and + adds only pressure. The mapping is then the right home for every weight + that has one, whatever the free-memory reading says. + """ + return current_platform.device_shares_host_memory() + + def host_copies_would_not_fit(weight_bytes: int) -> bool: """Whether copying `weight_bytes` into host memory would run the host out. @@ -186,7 +240,13 @@ class HostPinBudget: def __init__(self, available_bytes: int | None = None) -> None: if available_bytes is None: - available_bytes = host_memory_available_bytes() + if host_copies_are_redundant(): + # Nothing to pin for: on one pool the copy duplicates + # page-cache bytes the device can already read, and the mapped + # courier overlaps its transfers anyway. + available_bytes = 0 + else: + available_bytes = host_memory_available_bytes() self.available_bytes = available_bytes self.reserve_bytes = max( int(available_bytes * HOST_RESERVE_FRACTION), MIN_HOST_RESERVE_BYTES diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index 68d43809a..7a7600c94 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py @@ -1,7 +1,13 @@ import bisect +import ctypes +import ctypes.util +import mmap +import os import queue import re +import sys import threading +import time from collections.abc import Mapping, Sequence from contextlib import nullcontext from time import perf_counter @@ -29,6 +35,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( HostPinBudget, describe_host_memory, + host_copies_are_redundant, host_copies_would_not_fit, host_memory_available_bytes, module_weight_bytes, @@ -206,6 +213,354 @@ def _install_host_gather_hooks( module.register_forward_hook(_output_to_device) +_MADV_WILLNEED = 3 +_libc = None +if sys.platform == "linux": + try: + _libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True) + except OSError: + _libc = None +_PAGE = os.sysconf("SC_PAGESIZE") if hasattr(os, "sysconf") else 4096 + + +_WILLNEED_MIN_AVAILABLE = 4 << 30 # bytes of MemAvailable required to advise + + +def _willneed_headroom_ok(need_bytes: int) -> bool: + """Advised pages need somewhere to land, or the advice backfires. + + Field-measured on a 32 GB host with MemAvailable at 0.29 GiB: pages read + ahead were evicted before the courier reached them, so every byte was + read twice and effective throughput fell to 0.67x of the unadvised run + (0.85 vs 1.27 GB/s). Only advise when the kernel has real headroom to + keep the window resident until it is consumed. + """ + try: + with open("/proc/meminfo") as handle: + for line in handle: + if line.startswith("MemAvailable:"): + available = int(line.split()[1]) * 1024 + return available >= max(_WILLNEED_MIN_AVAILABLE, 2 * need_bytes) + except (OSError, ValueError): + pass + return False + + +def advise_willneed(tensors) -> int: + """Ask the kernel to read these mapped tensors' pages ahead, in bulk. + + The default readahead pipeline feeds the drive in read_ahead_kb-sized + beats (128 KB), which holds a fast NVMe to a quarter of its sequential + throughput on a host too small to cache the checkpoint. MADV_WILLNEED + schedules the whole range at once, so the disk read for the next layer + runs at drive speed while the current layer computes. Best-effort and + Linux-only: on any failure the normal fault path still works — and on a + host with no free headroom the advice is withheld entirely, because a + window that cannot stay resident until consumed is read twice. + """ + if _libc is None: + return 0 + tensors = list(tensors) + need = 0 + for tensor in tensors: + try: + need += tensor.untyped_storage().nbytes() + except Exception: + continue + if need == 0 or not _willneed_headroom_ok(need): + return 0 + advised = 0 + for tensor in tensors: + try: + storage = tensor.untyped_storage() + ptr = storage.data_ptr() + nbytes = storage.nbytes() + except Exception: + continue + if ptr == 0 or nbytes == 0: + continue + start = ptr & ~(_PAGE - 1) + length = (ptr + nbytes) - start + length = (length + _PAGE - 1) & ~(_PAGE - 1) + if ( + _libc.madvise( + ctypes.c_void_p(start), ctypes.c_size_t(length), _MADV_WILLNEED + ) + == 0 + ): + advised += 1 + return advised + + +_MADV_COLD = 20 # Linux 5.4+: deactivate the pages; reclaimed first under pressure +_MADV_PAGEOUT = 21 # Linux 5.4+: reclaim the pages now + + +_MADV_POPULATE_READ = 22 # Linux 5.14+: fault the range in, in one sequential pass +# Streamed layers faulted in concurrently on a cold pass. One sequential +# stream gets ~1 GiB/s from a GB10's NVMe. Within one checkpoint shard six +# streams measured 1.7 GiB/s and twelve 3.0 GiB/s, so the pool runs twelve. +MAPPED_POPULATE_AHEAD = 12 + + +def populate_mapped_source(tensors) -> int: + """Fault a layer's mapped pages in before a parallel copy reads them. + + A multi-threaded memcpy over an uncached mapping faults from many offsets + at once, which the kernel's readahead heuristics read as random access: + the drive is then fed 4 KiB at a time and a 1.1 GiB/s NVMe delivers a + fifth of that (a 125 s first denoise step on a GB10). One synchronous + MADV_POPULATE_READ per tensor keeps the read sequential and full-speed; + on an older kernel it fails with EINVAL and the WILLNEED path remains. + """ + if _libc is None: + return 0 + populated = 0 + for tensor in tensors: + try: + ptr = tensor.data_ptr() + nbytes = tensor.numel() * tensor.element_size() + except Exception: + continue + if ptr == 0 or nbytes == 0: + continue + start = ptr & ~(_PAGE - 1) + length = (ptr + nbytes) - start + length = (length + _PAGE - 1) & ~(_PAGE - 1) + if ( + _libc.madvise( + ctypes.c_void_p(start), ctypes.c_size_t(length), _MADV_POPULATE_READ + ) + == 0 + ): + populated += 1 + else: + advise_willneed([tensor]) + return populated + + +def _advise_mapped_source_cold(tensor: torch.Tensor, *, reclaim: bool = False) -> None: + """Tell the kernel a mapped tensor's file pages are cold once a copy holds them. + + MADV_COLD deactivates the pages without dropping them: under pressure the + kernel reclaims them ahead of anything hot, and otherwise keeps them. A + just-copied 45 GiB encoder otherwise looks like the hottest data on the + box and the kernel swaps idle anonymous memory instead -- measured on a + GB10 as 15 GiB of swap traffic and a wedged host. + + ``reclaim`` asks for MADV_PAGEOUT instead: the pages go now. That is for + a permanent materialization on a shared pool, where the device grows by + the same bytes the source pages hold and the driver satisfies device + allocations from free memory without waiting for cache reclaim -- a 57 GiB + DiT copied with MemFree near zero ended in NVRM out-of-memory twice. + """ + if not sys.platform.startswith("linux") or tensor.device.type != "cpu": + return + if _libc is None: + return + page = mmap.PAGESIZE + start = tensor.data_ptr() + end = start + tensor.numel() * tensor.element_size() + start -= start % page + if end <= start: + return + _libc.madvise( + ctypes.c_void_p(start), + ctypes.c_size_t(end - start), + _MADV_PAGEOUT if reclaim else _MADV_COLD, + ) + + +def _shared_pool_hosting( + totals: Dict[int, int], mapped: Dict[int, int] +) -> Dict[int, str]: + """Hosting when host and device draw from one pool. + + A mapped layer stays mapped: the device reads page-cache pages directly, so + a pinned or pageable copy would hold the same bytes twice. Only a layer with + no mapping at all -- an anonymous fused weight -- keeps a pageable copy. + """ + return { + layer_idx: "mapped" if mapped.get(layer_idx, 0) > 0 else "pageable" + for layer_idx in totals + } + + +_DIRECT_ALIGN = 4096 +_DIRECT_CHUNK = 64 << 20 +# Components smaller than this keep the page-cache path: their pages fit the +# cache next to a resident DiT, and a component re-streamed many times per +# request (the H3 video VAE: 36 layers, 4.5 GiB, ~200 passes per decode) +# must not go to the drive on every pass. +MAPPED_DIRECT_READ_MIN_BYTES = 8 * 1024**3 + + +class _DirectReader: + """Read a mapped tensor's bytes from its checkpoint file with O_DIRECT. + + Measured on a GB10: the NVMe delivers 9.9 GiB/s to an O_DIRECT reader, + but 1.1 GiB/s (one stream) to 3.7 GiB/s (twelve) through the page cache, + whose page allocation and reclaim are the wall on a shared pool. Reading + straight into the courier's pinned slot skips the cache entirely: no + pages to populate, evict or reclaim, and no cache footprint at all. + """ + + def __init__(self) -> None: + self._ranges: List[tuple[int, int, int, str]] = [] + self._starts: List[int] = [] + self._fds: Dict[str, int] = {} + self._located: Dict[int, Optional[tuple[str, int, int]]] = {} + self.refresh() + + def refresh(self) -> None: + ranges = [] + try: + with open("/proc/self/maps") as handle: + for line in handle: + fields = line.split() + if len(fields) < 6 or fields[5].startswith("["): + continue + start, end = (int(x, 16) for x in fields[0].split("-")) + ranges.append((start, end, int(fields[2], 16), fields[5])) + except OSError: + ranges = [] + ranges.sort() + self._ranges = ranges + self._starts = [r[0] for r in ranges] + self._located.clear() + + def locate(self, tensor: torch.Tensor) -> Optional[tuple[str, int, int]]: + """(path, file offset, nbytes) of the tensor's bytes, or None if unmapped.""" + ptr = tensor.data_ptr() + nbytes = tensor.numel() * tensor.element_size() + cached = self._located.get(ptr) + if cached is not None or ptr in self._located: + return cached + found = None + for attempt in range(2): + index = bisect.bisect_right(self._starts, ptr) - 1 + if index >= 0: + start, end, file_offset, path = self._ranges[index] + if start <= ptr and ptr + nbytes <= end: + found = (path, file_offset + (ptr - start), nbytes) + break + if attempt == 0: + self.refresh() + self._located[ptr] = found + return found + + def fd(self, path: str) -> int: + fd = self._fds.get(path) + if fd is None: + fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECT")) + self._fds[path] = fd + return fd + + def read_into( + self, buffer: memoryview, path: str, aligned_offset: int, span: int + ) -> None: + """Fill buffer[:span] with the file bytes at aligned_offset (both 4 KiB aligned).""" + fd = self.fd(path) + pos = 0 + while pos < span: + want = min(_DIRECT_CHUNK, span - pos) + got = os.preadv(fd, [buffer[pos : pos + want]], aligned_offset + pos) + if got <= 0: + # past the end of the file: the tail of the last aligned span + break + pos += got + + def close(self) -> None: + for fd in self._fds.values(): + try: + os.close(fd) + except OSError: + pass + self._fds.clear() + + +def _aligned_span(file_offset: int, nbytes: int) -> tuple[int, int]: + """(aligned start, span) covering [file_offset, file_offset + nbytes) at 4 KiB granularity.""" + start = file_offset & ~(_DIRECT_ALIGN - 1) + end = (file_offset + nbytes + _DIRECT_ALIGN - 1) & ~(_DIRECT_ALIGN - 1) + return start, end - start + + +class _MappedPopulator: + """Fault upcoming mapped layers in from several threads at once. + + On the pass that may find pages cold, the manager hands the next few + streamed layers to this pool so the drive stays saturated while the + courier stages and the model computes the current one. A layer is + populated once per request; ``reset`` starts the next request over. + """ + + def __init__(self, workers: int = MAPPED_POPULATE_AHEAD) -> None: + self._tasks: queue.Queue[Optional[tuple[int, list]]] = queue.Queue() + self._lock = threading.Condition() + self._pending: Set[int] = set() + self._done: Set[int] = set() + self.stats = {"populate_s": 0.0, "bytes": 0, "layers": 0} + self._threads = [ + threading.Thread(target=self._run, name=f"mapped-populate-{i}", daemon=True) + for i in range(max(1, workers)) + ] + for thread in self._threads: + thread.start() + + def submit(self, layer_idx: int, tensors) -> bool: + tensors = list(tensors) + if not tensors: + return False + with self._lock: + if layer_idx in self._pending or layer_idx in self._done: + return False + self._pending.add(layer_idx) + self._tasks.put((layer_idx, tensors)) + return True + + def wait(self, layer_idx: int) -> bool: + """Block until a submitted layer is populated; False if it never was.""" + with self._lock: + if layer_idx not in self._pending and layer_idx not in self._done: + return False + while layer_idx in self._pending: + self._lock.wait() + return layer_idx in self._done + + def reset(self) -> None: + with self._lock: + self._done.clear() + + def close(self) -> None: + for _ in self._threads: + self._tasks.put(None) + for thread in self._threads: + thread.join(timeout=5.0) + + def _run(self) -> None: + while True: + task = self._tasks.get() + if task is None: + return + layer_idx, tensors = task + started = time.perf_counter() + try: + populate_mapped_source(tensors) + except Exception: # advisory only; the courier reads the pages anyway + pass + with self._lock: + self.stats["populate_s"] += time.perf_counter() - started + self.stats["bytes"] += sum( + t.numel() * t.element_size() for t in tensors + ) + self.stats["layers"] += 1 + with self._lock: + self._pending.discard(layer_idx) + self._done.add(layer_idx) + self._lock.notify_all() + + class MappedLayerCourier: """Ships a mapped layer's weights to the device off the compute thread. @@ -232,24 +587,69 @@ class MappedLayerCourier: weight_metadata: Dict[int, Dict[str, Dict[str, Any]]], device: torch.device, pin_slots: bool, + cold_source: Optional[Callable[[int], bool]] = None, + populate_source: Optional[Callable[[int], bool]] = None, + await_populated: Optional[Callable[[int], bool]] = None, + direct_copy: bool = False, + direct_read: bool = False, ) -> None: self._mapped_cpu_weights = mapped_cpu_weights + # Read each layer's bytes from the checkpoint file with O_DIRECT into the + # pinned slot instead of copying them out of the page cache. + self.direct_read = bool(direct_read) and hasattr(os, "O_DIRECT") + self._reader: Optional[_DirectReader] = ( + _DirectReader() if self.direct_read else None + ) + self._slot_views: Dict[int, Any] = {} + # On a shared host/device pool the device reads host pages directly, so + # a layer goes mapping -> device in one copy and the pinned staging + # slots (2.7 GiB on the H3 encoder + VAE) are not allocated at all. + self._direct_copy = direct_copy + self.stats = { + "layers": 0, + "bytes": 0, + "slot_sync_s": 0.0, + "populate_wait_s": 0.0, + "populate_s": 0.0, + "memcpy_s": 0.0, + "h2d_issue_s": 0.0, + "direct_read_s": 0.0, + "direct_read_bytes": 0, + } + # Blocks until a populator thread has faulted the layer in; True if one + # did, so the courier does not fault the same range a second time. + self._await_populated = await_populated self._weight_metadata = weight_metadata self._device = device + # Whether a layer's file pages may go cold once its copy is staged: + # true for resident layers, never for layers re-read every step. + self._cold_source = cold_source + # Whether to fault the layer in sequentially before the staging copy: + # true on a request's first pass over the layers, when pages may be cold. + self._populate_source = populate_source slot_bytes = max( ( - sum(t.numel() * t.element_size() for t in weights.values()) + sum( + t.numel() * t.element_size() + + (2 * _DIRECT_ALIGN if self.direct_read else 0) + for t in weights.values() + ) for weights in mapped_cpu_weights.values() if weights ), default=0, ) + slot_bytes = (slot_bytes + _DIRECT_ALIGN - 1) & ~(_DIRECT_ALIGN - 1) if slot_bytes <= 0: raise ValueError("no mapped weights to ship") - self._slots = [ - torch.empty(slot_bytes, dtype=torch.uint8, pin_memory=pin_slots) - for _ in range(self._NUM_SLOTS) - ] + self._slots = ( + [] + if direct_copy + else [ + torch.empty(slot_bytes, dtype=torch.uint8, pin_memory=pin_slots) + for _ in range(self._NUM_SLOTS) + ] + ) self._slot_events: List[Optional[Any]] = [None] * self._NUM_SLOTS self._stream = torch.get_device_module().Stream() self._tasks: queue.Queue[Optional[int]] = queue.Queue() @@ -292,6 +692,8 @@ class MappedLayerCourier: def close(self) -> None: self._tasks.put(None) self._thread.join(timeout=5.0) + if self._reader is not None: + self._reader.close() def _run(self) -> None: slot_turn = 0 @@ -313,37 +715,129 @@ class MappedLayerCourier: return def _ship(self, layer_idx: int, slot_turn: int): - slot = self._slots[slot_turn] - previous = self._slot_events[slot_turn] - if previous is not None: - # the previous transfer through this slot must land before reuse - previous.synchronize() + stats = self.stats tensors: Dict[str, torch.Tensor] = {} - offset = 0 + if not self._direct_copy: + slot = self._slots[slot_turn] + previous = self._slot_events[slot_turn] + if previous is not None: + # the previous transfer through this slot must land before reuse + started = time.perf_counter() + previous.synchronize() + stats["slot_sync_s"] += time.perf_counter() - started + if ( + not self.direct_read + and self._populate_source is not None + and self._populate_source(layer_idx) + ): + started = time.perf_counter() + if self._await_populated is not None and self._await_populated(layer_idx): + stats["populate_wait_s"] += time.perf_counter() - started + else: + populate_mapped_source(self._mapped_cpu_weights[layer_idx].values()) + stats["populate_s"] += time.perf_counter() - started + layer_bytes = 0 with torch.inference_mode(False), torch.no_grad(): - staged = [] - for name, cpu_tensor in self._mapped_cpu_weights[layer_idx].items(): - width = cpu_tensor.element_size() - if offset % width: - offset += width - (offset % width) - start = offset // width - window = slot.view(cpu_tensor.dtype)[ - start : start + cpu_tensor.numel() - ].view(cpu_tensor.shape) - window.copy_(cpu_tensor) - offset += cpu_tensor.numel() * width - staged.append((name, window)) event = torch.get_device_module().Event() - with torch.get_device_module().stream(self._stream): - for name, window in staged: - meta = self._weight_metadata[layer_idx][name] - gpu_tensor = torch.empty( - meta["shape"], dtype=meta["dtype"], device=self._device + if self._direct_copy: + started = time.perf_counter() + with torch.get_device_module().stream(self._stream): + for name, cpu_tensor in self._mapped_cpu_weights[layer_idx].items(): + meta = self._weight_metadata[layer_idx][name] + gpu_tensor = torch.empty( + meta["shape"], dtype=meta["dtype"], device=self._device + ) + gpu_tensor.copy_(cpu_tensor, non_blocking=True) + layer_bytes += cpu_tensor.numel() * cpu_tensor.element_size() + if self._cold_source is not None and self._cold_source( + layer_idx + ): + _advise_mapped_source_cold(cpu_tensor, reclaim=True) + tensors[name] = gpu_tensor + event.record(self._stream) + stats["h2d_issue_s"] += time.perf_counter() - started + else: + offset = 0 + staged = [] + started = time.perf_counter() + slot_bytes_view = None + if self.direct_read: + slot_bytes_view = self._slot_views.get(slot_turn) + if slot_bytes_view is None: + slot_bytes_view = memoryview(slot.numpy()) + self._slot_views[slot_turn] = slot_bytes_view + for name, cpu_tensor in self._mapped_cpu_weights[layer_idx].items(): + width = cpu_tensor.element_size() + nbytes = cpu_tensor.numel() * width + located = ( + self._reader.locate(cpu_tensor) if self.direct_read else None ) - gpu_tensor.copy_(window, non_blocking=True) - tensors[name] = gpu_tensor - event.record(self._stream) - self._slot_events[slot_turn] = event + if located is not None: + path, file_offset, _ = located + aligned_start, span = _aligned_span(file_offset, nbytes) + skew = file_offset - aligned_start + if offset % _DIRECT_ALIGN: + offset += _DIRECT_ALIGN - (offset % _DIRECT_ALIGN) + if skew % width == 0 and offset + span <= slot.numel(): + read_started = time.perf_counter() + try: + self._reader.read_into( + slot_bytes_view[offset : offset + span], + path, + aligned_start, + span, + ) + except OSError as exc: + logger.warning( + "Layerwise offload: O_DIRECT read failed (%s); " + "mapped layers return to the page-cache copy.", + exc, + ) + self.direct_read = False + located = None + else: + stats["direct_read_s"] += ( + time.perf_counter() - read_started + ) + stats["direct_read_bytes"] += span + window = ( + slot[offset + skew : offset + skew + nbytes] + .view(cpu_tensor.dtype) + .view(cpu_tensor.shape) + ) + offset += span + else: + located = None + if located is None: + if offset % width: + offset += width - (offset % width) + start = offset // width + window = slot.view(cpu_tensor.dtype)[ + start : start + cpu_tensor.numel() + ].view(cpu_tensor.shape) + window.copy_(cpu_tensor) + if self._cold_source is not None and self._cold_source( + layer_idx + ): + _advise_mapped_source_cold(cpu_tensor, reclaim=True) + offset += nbytes + layer_bytes += nbytes + staged.append((name, window)) + stats["memcpy_s"] += time.perf_counter() - started + started = time.perf_counter() + with torch.get_device_module().stream(self._stream): + for name, window in staged: + meta = self._weight_metadata[layer_idx][name] + gpu_tensor = torch.empty( + meta["shape"], dtype=meta["dtype"], device=self._device + ) + gpu_tensor.copy_(window, non_blocking=True) + tensors[name] = gpu_tensor + event.record(self._stream) + stats["h2d_issue_s"] += time.perf_counter() - started + self._slot_events[slot_turn] = event + stats["layers"] += 1 + stats["bytes"] += layer_bytes return event, tensors @@ -414,6 +908,14 @@ class LayerwiseOffloadManager: # Armed on the first denoise forward, so that the load-time prefetch below # does not pin the whole resident set before the DiT is the active component. self._residency_active = False + # True while load_all_layers materializes every layer for a resident + # placement; every mapped source is then read exactly once. + self._materializing_all = False + # True from a request's start until its last layer has run once: the + # pass in which a mapped layer's pages may not be in the page cache. + self._first_pass = True + self._mapped_populator: Optional[_MappedPopulator] = None + self._debug_collect_wait_s = 0.0 # True once _initialize builds the CPU buffers; unlike `enabled` it # never flips back, so disable_offload/enable_offload can toggle # `enabled` without losing track of which managers can be re-armed. @@ -614,6 +1116,19 @@ class LayerwiseOffloadManager: buys a whole layer's worth of per-step overlap. """ totals, mapped = self._layer_byte_totals(layer_groups) + if host_copies_are_redundant(): + hosting = _shared_pool_hosting(totals, mapped) + logger.info( + "Layerwise offload: %s keeps %d of %d layers on the checkpoint " + "mapping (host and device share one memory pool, so a pinned or " + "pageable copy would hold the same bytes twice); %d layers " + "without a mapping stay pageable.", + self._pin_component_name, + sum(1 for where in hosting.values() if where == "mapped"), + len(totals), + sum(1 for where in hosting.values() if where == "pageable"), + ) + return hosting pinned_bytes = 0 hosting: Dict[int, str] = {} pin_order: List[int] = [] @@ -956,6 +1471,134 @@ class LayerwiseOffloadManager: if layer_idx not in retain: self.release_layer(layer_idx) + def advise_mapped_pages_cold(self, *, room_bytes: Optional[int] = None) -> int: + """Hand this stage's mapped pages back in the order that keeps the next + request fast. Returns the bytes paged out. + + With host and device in one pool the page cache cannot always hold the + whole request cycle. Plain LRU then evicts, at each phase boundary, + exactly the pages the next phase needs: a cyclic scan just larger + than the cache misses everywhere (measured on a GB10: ~100 GiB re-read + per request). Deactivating a component's pages (MADV_COLD) once its + stage is over fixes the cross-component case. + + When the component's own stream is larger than the room the cache will + have for it (``room_bytes``), the same pathology happens inside the + stage: faulting layer 40 in evicts layer 0, which the next request + reads first. Measured on a GB10 with a 45 GiB encoder and ~30 GiB of + room, every page fault reclaimed synchronously and the stage took + 38 s at 1.2 GiB/s. So the head of the stream -- as many layers as do + not fit -- is paged out now, deterministically, and the tail is left + cold: the next request reads the head from disk at full parallel + speed into free pages and hits the cache for everything else. + """ + order = [ + idx for idx in self._streamed_order if self._mapped_cpu_weights.get(idx) + ] + order += [ + idx + for idx in self._mapped_cpu_weights + if idx not in set(order) and self._mapped_cpu_weights.get(idx) + ] + layer_bytes = { + idx: sum( + t.numel() * t.element_size() + for t in self._mapped_cpu_weights[idx].values() + ) + for idx in order + } + total = sum(layer_bytes.values()) + excess = 0 + if room_bytes is not None and total > room_bytes: + excess = total - room_bytes + paged_out = 0 + for idx in order: + reclaim = paged_out < excess + for tensor in self._mapped_cpu_weights[idx].values(): + _advise_mapped_source_cold(tensor, reclaim=reclaim) + if reclaim: + paged_out += layer_bytes[idx] + return paged_out + + def _ensure_mapped_populator(self) -> Optional[_MappedPopulator]: + if self._mapped_populator is None and _libc is not None: + self._mapped_populator = _MappedPopulator() + return self._mapped_populator + + def _await_mapped_populated(self, layer_idx: int) -> bool: + populator = self._mapped_populator + return populator is not None and populator.wait(layer_idx) + + def _populate_ahead(self, layer_idx: int) -> None: + """On a cold pass, fault the next streamed layers in from parallel threads.""" + populator = self._ensure_mapped_populator() + if populator is None: + return + for ahead in self._next_streamed(after=layer_idx, count=MAPPED_POPULATE_AHEAD): + if ahead in self._gpu_layers or ahead in self._courier_inflight: + continue + populator.submit(ahead, self._mapped_cpu_weights.get(ahead, {}).values()) + + def _log_debug_timing(self) -> None: + """Debug: where this stage's layer traffic spent its time.""" + if not envs.SGLANG_DIFFUSION_DEBUG_LAYERWISE_TIMING: + return + courier = self._mapped_courier + populator = self._mapped_populator + if courier is None or not courier.stats["layers"]: + self._debug_collect_wait_s = 0.0 + return + cs = courier.stats + ps = ( + populator.stats + if populator is not None + else {"populate_s": 0.0, "bytes": 0, "layers": 0} + ) + logger.info( + "Layerwise timing %s: layers=%d bytes=%.2fGiB direct=%s | courier: slot_sync=%.2fs " + "populate_wait=%.2fs populate=%.2fs memcpy=%.2fs (direct_read=%.2fs %.2fGiB) " + "h2d_issue=%.2fs | populator: " + "layers=%d bytes=%.2fGiB busy=%.2fs | compute thread collect wait=%.2fs", + self.layers_attr_str, + cs["layers"], + cs["bytes"] / (1024**3), + courier._direct_copy, + cs["slot_sync_s"], + cs["populate_wait_s"], + cs["populate_s"], + cs["memcpy_s"], + cs["direct_read_s"], + cs["direct_read_bytes"] / (1024**3), + cs["h2d_issue_s"], + ps["layers"], + ps["bytes"] / (1024**3), + ps["populate_s"], + self._debug_collect_wait_s, + ) + for key in cs: + cs[key] = 0.0 if isinstance(cs[key], float) else 0 + if populator is not None: + with populator._lock: + for key in ps: + ps[key] = 0.0 if isinstance(ps[key], float) else 0 + self._debug_collect_wait_s = 0.0 + + def _mapped_source_may_be_cold(self, layer_idx: int) -> bool: + """Whether this copy is the first read of the layer in this request.""" + return self._first_pass or self._materializing_all + + def _mapped_source_is_cold(self, layer_idx: int) -> bool: + """Whether a mapped layer's file pages may go once its device copy lands. + + Only while every layer is being materialized for a permanent resident + placement: those pages are not read again until a demotion. A + stage-scoped resident set is re-armed from the same pages on the next + request, and a streamed layer is re-read every step; marking either + cold made the kernel evict exactly what the next request needed + (measured on a GB10 as a 60 s first denoise step re-reading 23 GiB). + """ + return self._materializing_all + @torch.compiler.disable def _activate_residency(self) -> None: """Arm the resident set on the first denoise forward. The pinning itself is @@ -1021,6 +1664,21 @@ class LayerwiseOffloadManager: if courier is not None and courier.submit(layer_idx): self._courier_inflight.add(layer_idx) ship_mapped = True + if ( + not envs.SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED + and not courier.direct_read + ): + # Schedule the disk read for this layer's pages now, in + # one bulk request, so it overlaps the previous layer's + # compute instead of trickling in at fault-time beats. + advise_willneed(self._mapped_cpu_weights[layer_idx].values()) + if self._first_pass: + # On the pass that may find pages cold, keep the drive + # saturated several layers ahead from parallel threads: + # one sequential stream idles the NVMe at ~1 GiB/s + # while a layer is staged and computed; twelve streams + # within one shard measured ~3 GiB/s on a GB10. + self._populate_ahead(layer_idx) # create gpu buffer and load from CPU buffer gpu_buffers: Dict[torch.dtype, torch.Tensor] = {} @@ -1050,10 +1708,16 @@ class LayerwiseOffloadManager: # on the compute thread rather than ahead of it, and a page # the kernel has reclaimed is faulted back in here. cpu_tensor = self._mapped_cpu_weights[layer_idx][name] + if not envs.SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED: + # A blocking read on this thread: fault the range in + # sequentially so a cold cache fills at drive speed. + populate_mapped_source([cpu_tensor]) gpu_tensor = torch.empty( meta["shape"], dtype=meta["dtype"], device=self.device ) gpu_tensor.copy_(cpu_tensor, non_blocking=False) + if self._mapped_source_is_cold(layer_idx): + _advise_mapped_source_cold(cpu_tensor, reclaim=True) target.data = self._wrap_for_target(target, gpu_tensor) continue @@ -1107,6 +1771,19 @@ class LayerwiseOffloadManager: weight_metadata=self._weight_metadata, device=self.device, pin_slots=current_platform.is_cuda(), + # Measured on a GB10: a copy_ straight from file-backed pages ran + # at 0.67 GiB/s (the driver stages pageable sources page by + # page) and the process's anonymous memory grew past 100 GiB; + # the pinned slots stay even on a shared pool. + direct_copy=False, + direct_read=( + host_copies_are_redundant() + and not envs.SGLANG_DIFFUSION_DISABLE_MAPPED_DIRECT_READ + and self._mapped_bytes >= MAPPED_DIRECT_READ_MIN_BYTES + ), + cold_source=self._mapped_source_is_cold, + populate_source=self._mapped_source_may_be_cold, + await_populated=self._await_mapped_populated, ) logger.info( "Layerwise offload: %s ships mapped layers through a courier " @@ -1128,6 +1805,7 @@ class LayerwiseOffloadManager: def _collect_mapped_layer(self, layer_idx: int) -> None: """Bind a shipped layer's tensors on the compute thread.""" courier = self._mapped_courier + started = time.perf_counter() try: event, tensors = courier.collect(layer_idx) except BaseException as exc: @@ -1143,6 +1821,10 @@ class LayerwiseOffloadManager: self._courier_inflight.discard(layer_idx) self.prefetch_layer(layer_idx, non_blocking=False) return + # debug counter; a manager built without __init__ (tests) has none yet + self._debug_collect_wait_s = getattr(self, "_debug_collect_wait_s", 0.0) + ( + time.perf_counter() - started + ) compute_stream = torch.get_device_module().current_stream() compute_stream.wait_event(event) with torch.inference_mode(False), torch.no_grad(): @@ -1195,6 +1877,9 @@ class LayerwiseOffloadManager: def release_all(self) -> None: """Release every layer, including the resident ones: this ends the denoise stage that the resident set is scoped to.""" + self._log_debug_timing() + if self._mapped_populator is not None: + self._mapped_populator.reset() if not self.enabled or self.device is None: return if self.copy_stream is not None: @@ -1207,6 +1892,9 @@ class LayerwiseOffloadManager: for layer_idx in list(self._gpu_layers): self.release_layer(layer_idx, force=True) + # The next use starts a new request; its first pass over the layers may + # find their pages evicted and is the one worth faulting in sequentially. + self._first_pass = True @torch.compiler.disable def load_all_layers(self) -> None: @@ -1216,9 +1904,54 @@ class LayerwiseOffloadManager: if self.copy_stream is not None: torch.get_device_module().current_stream().wait_stream(self.copy_stream) - for layer_idx in range(self.num_layers): - if layer_idx not in self._gpu_layers: - self.prefetch_layer(layer_idx, non_blocking=False) + self._materializing_all = True + try: + for layer_idx in range(self.num_layers): + if layer_idx not in self._gpu_layers: + # Anonymous host stores can fill the copy stream without a + # per-layer host wait. Checkpoint mappings still use the + # synchronous path: the mapped courier has a bounded slot + # ring intended to overlap one forward, not materialize a + # whole model at once. + self.prefetch_layer( + layer_idx, + non_blocking=not bool(self._mapped_cpu_weights.get(layer_idx)), + ) + finally: + self._materializing_all = False + if self.copy_stream is not None: + torch.get_device_module().current_stream().wait_stream(self.copy_stream) + + def release_host_stores(self) -> None: + """Drop rollback stores after a resident placement is validated. + + The real device tensors must already be materialized and the manager + disabled. Repacking pinned stores as pageable here would copy the full + checkpoint for data that will never be streamed again. + """ + if self.enabled: + raise RuntimeError("cannot release host stores while offload is enabled") + if self._mapped_courier is not None: + self._mapped_courier.close() + self._mapped_courier = None + if self._mapped_populator is not None: + self._mapped_populator.close() + self._mapped_populator = None + if self._courier_inflight: + raise RuntimeError( + "cannot release host stores with mapped copies in flight" + ) + + self._pin_budget.release(self.pinned_host_weight_bytes()) + self._consolidated_cpu_weights.clear() + self._strided_cpu_weights.clear() + self._mapped_cpu_weights.clear() + self._mps_cpu_weights.clear() + self._weight_metadata.clear() + self._layer_hosting.clear() + self._prefetch_events.clear() + self._mapped_bytes = 0 + self._configured = False @torch.compiler.disable def sync_layer_to_cpu(self, layer_idx: int) -> None: @@ -1842,6 +2575,12 @@ class LayerwiseOffloadableModuleMixin: policies = ", ".join(sorted({manager.residency_policy for manager in managers})) total_layers = sum(manager.num_layers for manager in managers) resident_layers = sum(manager.resident_layers for manager in managers) + if envs.SGLANG_DIFFUSION_DEBUG_HOST_MEMORY: + from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_breakdown import ( + log_anon_vmas, + ) + + log_anon_vmas(f"layerwise offload ready for {component_name}") logger.info( "Layerwise offload ready for %s in %.2fs: groups=%d, layers=%d, " "prefetch/group=%s, resident=%d/%d, policy=%s", diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 844ae4006..6acf931e6 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -45,6 +45,12 @@ pynvml = import_pynvml() # type: ignore[no-untyped-call] torch.backends.cuda.enable_cudnn_sdp(False) +@lru_cache(maxsize=None) +def _device_is_integrated(device_index: int) -> bool: + # A static device property, asked on every planner cost evaluation. + return bool(torch.cuda.get_device_properties(device_index).is_integrated) + + def device_id_to_physical_device_id(device_id: int) -> int: if "CUDA_VISIBLE_DEVICES" in os.environ: device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",") @@ -628,6 +634,15 @@ class CudaPlatformBase(Platform): return free_gpu_memory / (1 << 30) + @classmethod + def device_shares_host_memory(cls) -> bool: + if not torch.cuda.is_available(): + return False + try: + return _device_is_integrated(torch.cuda.current_device()) + except (RuntimeError, AssertionError): + return False + @classmethod def _resolve_default_attn_backend(cls) -> AttentionBackendEnum: if cls.is_sm120(): diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 04eb9baf2..8e24d7c39 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -434,6 +434,16 @@ class Platform: """Whether automatic DiT layerwise offload is enabled on this platform.""" return True + @classmethod + def device_shares_host_memory(cls) -> bool: + """Whether the accelerator draws from the same physical pool as the host. + + On such a part (DGX Spark's GB10, Jetson) a device allocation is host + memory the kernel no longer has, and a host copy of a mapped weight is + a second copy of bytes the page cache already holds. + """ + return False + @classmethod def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module: """Apply platform-specific optimizations to VAE after loading.""" diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index 34abb57e6..0b8a8a038 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -4,6 +4,7 @@ ServerArgsAutoTuner tunes the ServerArgs based on the desired performance mode from __future__ import annotations +import os from typing import TYPE_CHECKING from sglang.multimodal_gen import envs @@ -88,6 +89,28 @@ IMAGE_GEN_KEEP_RESIDENT_MIN_AVAILABLE_GB = 45.0 DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB = 120.0 +# torch's CPU allocator (mimalloc since 2.13) keeps freed pages in its arenas +# and backs them with transparent huge pages, so the fused-weight copies the +# loader frees once a component is promoted stayed resident: 18.2 GiB of +# anonymous memory on a GB10 after the DiT went resident, 5.0 GiB with these. +# Read at process start, so they are set for the workers to inherit; a torch +# without mimalloc ignores them. +SHARED_POOL_CPU_ALLOCATOR_DEFAULTS = { + "MIMALLOC_PURGE_DELAY": "0", + "MIMALLOC_ALLOW_LARGE_OS_PAGES": "0", +} + + +def apply_shared_pool_cpu_allocator_defaults(environ) -> list[str]: + """Set the CPU allocator defaults not already chosen; return the names set.""" + applied = [] + for name, value in SHARED_POOL_CPU_ALLOCATOR_DEFAULTS.items(): + if name not in environ: + environ[name] = value + applied.append(name) + return applied + + class ServerArgsAutoTuner: """Auto-tunes the server-arg for the given performance-mode, based on practical deployment experience with different model architectures""" @@ -456,6 +479,104 @@ class ServerArgsAutoTuner: args.text_encoder_cpu_offload = False if args.image_encoder_cpu_offload is None: args.image_encoder_cpu_offload = False + if ( + args.pin_cpu_memory + and not args.is_arg_explicitly_set("pin_cpu_memory") + and current_platform.device_shares_host_memory() + ): + # The device reads host pages directly on a shared pool, so a + # pinned copy of a mapped weight is the same bytes held twice. + args.pin_cpu_memory = False + logger.info( + "Host and device share one memory pool: pinned host weight " + "copies are disabled (pass --pin-cpu-memory true to override)." + ) + if ( + current_platform.device_shares_host_memory() + and "PYTORCH_CUDA_ALLOC_CONF" not in os.environ + ): + # Every byte the caching allocator keeps reserved is a byte the + # page cache -- the home of every mapped weight here -- cannot + # hold. Measured on a GB10: ~30 GiB of reserved-but-idle segments + # forced the encoder and the DiT to take turns being re-read from + # disk. Expandable segments let the reserve follow the live peak. + os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + logger.info( + "Host and device share one memory pool: PYTORCH_CUDA_ALLOC_CONF=" + "expandable_segments:True so the allocator's reserve does not " + "crowd out the page cache." + ) + if current_platform.device_shares_host_memory(): + applied = apply_shared_pool_cpu_allocator_defaults(os.environ) + if applied: + logger.info( + "Host and device share one memory pool: %s so the CPU " + "allocator returns freed weight copies to the pool.", + " ".join(f"{name}={os.environ[name]}" for name in applied), + ) + if current_platform.device_shares_host_memory(): + try: + import psutil + + swap_total = psutil.swap_memory().total + except Exception: + swap_total = 0 + try: + # A cgroup with swap disabled (memory.swap.max = 0) protects + # this process whatever the host has mounted. + with open("/sys/fs/cgroup/memory.swap.max") as handle: + if handle.read().strip() == "0": + swap_total = 0 + except OSError: + pass + try: + with open("/sys/fs/cgroup/memory.max") as handle: + uncapped = handle.read().strip() == "max" + except OSError: + uncapped = True + if uncapped: + # The driver takes device memory from free pages and does not + # wait for the kernel to reclaim page cache: with the cache + # full and MemFree near zero, device growth fails outright + # (NVRM out-of-memory on a GB10, three runs). A cgroup limit a + # little under physical memory makes the kernel reclaim this + # process's cache ahead of its own allocations. + logger.warning( + "Host and device share one memory pool and this process has " + "no cgroup memory limit: device allocations may fail while " + "the page cache holds the free memory. Run with a limit a few " + "GiB under physical memory (for example docker --memory)." + ) + if swap_total > 0: + # Under page-cache pressure the kernel prefers swapping idle + # anonymous memory -- here the DiT's fused weight copies -- + # over dropping cache, and every denoise step then swaps them + # back in. Measured on a GB10: 128 s first steps and a 54 s + # text encoder with 143 GiB of swap enabled. + logger.warning( + "Host and device share one memory pool and swap is enabled " + "(%.0f GiB): the kernel may swap out weight copies under " + "page-cache pressure. Run with swap off for this process " + "(container --memory-swap equal to --memory, or " + "vm.swappiness=0).", + swap_total / 1024**3, + ) + if args.dit_cpu_offload or args.text_encoder_cpu_offload: + # Whole-component offload holds a component twice while it + # moves: the device copy plus a host copy the size of the + # component. On a shared pool both come out of the same + # memory. Measured on a GB10: a 57 GiB DiT moving back to + # the host at the end of a denoise stage exhausted the pool. + logger.warning( + "Host and device share one memory pool and whole-component " + "CPU offload is enabled (dit_cpu_offload=%s, " + "text_encoder_cpu_offload=%s): moving a component holds it " + "twice while it moves. Prefer layerwise offload, where " + "residency is armed layer by layer from the checkpoint " + "mapping.", + bool(args.dit_cpu_offload), + bool(args.text_encoder_cpu_offload), + ) def _normalize_performance_mode(self) -> str: args = self.server_args diff --git a/python/sglang/multimodal_gen/test/unit/test_host_spill.py b/python/sglang/multimodal_gen/test/unit/test_host_spill.py new file mode 100644 index 000000000..c5ecf6d72 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_host_spill.py @@ -0,0 +1,104 @@ +import os +import sys + +import pytest +import torch + +from sglang.multimodal_gen.runtime.loader import host_spill as host_spill_module +from sglang.multimodal_gen.runtime.loader.host_spill import ( + HostSpill, + checkpoint_fingerprint, +) +from sglang.multimodal_gen.runtime.loader.utils import ( + MappedRegions, + hf_to_custom_state_dict, +) + + +@pytest.fixture(autouse=True) +def _small_spill_threshold(monkeypatch): + monkeypatch.setattr(host_spill_module, "MIN_SPILL_BYTES", 0) + monkeypatch.setattr(host_spill_module, "SPILL_DISK_RESERVE_BYTES", 0) + + +def test_spilled_tensor_is_file_backed_and_reused_after_sealing(tmp_path): + spill = HostSpill(tmp_path, "ckpt") + shape = torch.Size([4, 8]) + tensor, filled = spill.tensor("blocks.0.qkv", shape, torch.bfloat16) + assert not filled + tensor.copy_(torch.arange(32, dtype=torch.bfloat16).view(4, 8)) + if sys.platform == "linux": + assert MappedRegions().holds(tensor) + # unsealed: the next start must not trust it + again, filled = HostSpill(tmp_path, "ckpt").tensor( + "blocks.0.qkv", shape, torch.bfloat16 + ) + assert not filled + spill.seal("blocks.0.qkv", shape, torch.bfloat16) + reused, filled = HostSpill(tmp_path, "ckpt").tensor( + "blocks.0.qkv", shape, torch.bfloat16 + ) + assert filled + assert torch.equal(reused.float(), torch.arange(32, dtype=torch.float32).view(4, 8)) + # a different dtype or shape is a different file + other, filled = HostSpill(tmp_path, "ckpt").tensor( + "blocks.0.qkv", shape, torch.float16 + ) + assert not filled + del again, other + + +def test_spill_disables_itself_when_the_disk_is_full(tmp_path, monkeypatch): + spill = HostSpill(tmp_path, "ckpt") + + class _Usage: + free = 0 + + monkeypatch.setattr(host_spill_module.shutil, "disk_usage", lambda _p: _Usage()) + monkeypatch.setattr(host_spill_module, "SPILL_DISK_RESERVE_BYTES", 1 << 30) + assert spill.tensor("w", torch.Size([2, 2]), torch.float32) is None + assert spill.tensor("w2", torch.Size([2, 2]), torch.float32) is None + assert spill.count_written == 0 + + +@pytest.mark.parametrize("strict", [False, True]) +def test_fused_weights_are_concatenated_into_the_provided_tensor(tmp_path, strict): + spill = HostSpill(tmp_path, "ckpt") + q = torch.full((2, 3), 1.0) + k = torch.full((2, 3), 2.0) + v = torch.full((2, 3), 3.0) + + def mapping(name): + prefix, _, which = name.rpartition(".") + return f"{prefix}.qkv", {"q": 0, "k": 1, "v": 2}[which], 3 + + weights = [("blocks.0.q", q), ("blocks.0.k", k), ("blocks.0.v", v)] + merged, _ = hf_to_custom_state_dict( + iter(weights), mapping, fused_tensor_factory=spill.tensor, strict=strict + ) + fused = merged["blocks.0.qkv"] + assert torch.equal(fused, torch.cat([q, k, v], dim=0)) + if sys.platform == "linux": + assert MappedRegions().holds(fused) + assert spill.count_written == 1 + spill.seal("blocks.0.qkv", fused.shape, fused.dtype) + + # the next start reuses the sealed file without reading the pieces + reused_spill = HostSpill(tmp_path, "ckpt") + merged_again, _ = hf_to_custom_state_dict( + iter([(n, torch.zeros_like(t)) for n, t in weights]), + mapping, + fused_tensor_factory=reused_spill.tensor, + strict=strict, + ) + assert torch.equal(merged_again["blocks.0.qkv"], torch.cat([q, k, v], dim=0)) + assert reused_spill.count_reused == 1 + + +def test_fingerprint_changes_with_the_checkpoint_files(tmp_path): + shard = tmp_path / "model-00001-of-00002.safetensors" + shard.write_bytes(b"a" * 16) + before = checkpoint_fingerprint([str(tmp_path)]) + shard.write_bytes(b"b" * 32) + os.utime(shard, ns=(1, 1)) + assert checkpoint_fingerprint([str(tmp_path)]) != before diff --git a/python/sglang/multimodal_gen/test/unit/test_mapped_willneed.py b/python/sglang/multimodal_gen/test/unit/test_mapped_willneed.py new file mode 100644 index 000000000..f75ac647e --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_mapped_willneed.py @@ -0,0 +1,93 @@ +"""advise_willneed hands the kernel page-aligned ranges and never raises. + +What matters: the madvise call receives a page-aligned start and a length +that covers the tensor's storage, odd offsets round outward, and failures +(no libc, bad pointers) degrade to zero advice instead of an exception. +""" + +import pytest +import torch + +from sglang.multimodal_gen.runtime.managers.memory_managers import ( + layerwise_offload as lo, +) + + +class _RecordingLibc: + def __init__(self, ret=0): + self.calls = [] + self.ret = ret + + def madvise(self, addr, length, advice): + self.calls.append((addr.value, length.value, advice)) + return self.ret + + +@pytest.fixture() +def libc(monkeypatch): + fake = _RecordingLibc() + monkeypatch.setattr(lo, "_libc", fake) + monkeypatch.setattr(lo, "_willneed_headroom_ok", lambda need: True) + return fake + + +def test_ranges_are_page_aligned_and_cover_the_storage(libc): + t = torch.zeros(1024, dtype=torch.float32) + advised = lo.advise_willneed([t]) + + assert advised == 1 + ((addr, length, advice),) = libc.calls + page = lo._PAGE + assert advice == lo._MADV_WILLNEED + assert addr % page == 0 + ptr = t.untyped_storage().data_ptr() + nbytes = t.untyped_storage().nbytes() + assert addr <= ptr + assert addr + length >= ptr + nbytes + assert length % page == 0 + + +def test_a_failing_madvise_counts_nothing(monkeypatch): + monkeypatch.setattr(lo, "_libc", _RecordingLibc(ret=-1)) + assert lo.advise_willneed([torch.zeros(16)]) == 0 + + +def test_no_libc_is_a_quiet_noop(monkeypatch): + monkeypatch.setattr(lo, "_libc", None) + assert lo.advise_willneed([torch.zeros(16)]) == 0 + + +def test_empty_and_broken_tensors_are_skipped(libc): + class Broken: + def untyped_storage(self): + raise RuntimeError("no storage") + + assert lo.advise_willneed([Broken(), torch.empty(0)]) == 0 + assert libc.calls == [] + + +def test_no_headroom_withholds_the_advice(monkeypatch): + fake = _RecordingLibc() + monkeypatch.setattr(lo, "_libc", fake) + monkeypatch.setattr(lo, "_willneed_headroom_ok", lambda need: False) + assert lo.advise_willneed([torch.zeros(1024)]) == 0 + assert fake.calls == [] + + +def test_headroom_reads_memavailable(monkeypatch, tmp_path): + meminfo = tmp_path / "meminfo" + + real_open = open + + def fake_open(path, *a, **k): + if path == "/proc/meminfo": + return real_open(meminfo, *a, **k) + return real_open(path, *a, **k) + + monkeypatch.setattr("builtins.open", fake_open) + + meminfo.write_text("MemTotal: 32 kB\nMemAvailable: 16777216 kB\n") # 16 GiB + assert lo._willneed_headroom_ok(1 << 30) + + meminfo.write_text("MemTotal: 32 kB\nMemAvailable: 524288 kB\n") # 0.5 GiB + assert not lo._willneed_headroom_ok(1 << 30) diff --git a/python/sglang/multimodal_gen/test/unit/test_readonly_safetensors.py b/python/sglang/multimodal_gen/test/unit/test_readonly_safetensors.py new file mode 100644 index 000000000..05f7c1b99 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_readonly_safetensors.py @@ -0,0 +1,82 @@ +import sys + +import pytest +import torch +from safetensors.torch import load_file, save_file + +from sglang.multimodal_gen.runtime.loader.readonly_safetensors import ( + iter_safetensors_readonly, + load_safetensors_readonly, + safetensors_keys, +) + + +def _write(tmp_path): + tensors = { + "a.weight": torch.randn(64, 32, dtype=torch.bfloat16), + "b.bias": torch.arange(17, dtype=torch.float32), + "c.empty": torch.empty(0, 4, dtype=torch.float16), + "d.flag": torch.tensor([True, False]), + "e.int": torch.arange(6, dtype=torch.int64).view(2, 3), + } + path = tmp_path / "model.safetensors" + save_file(tensors, str(path)) + return path, tensors + + +def test_readonly_load_matches_safetensors(tmp_path): + path, tensors = _write(tmp_path) + ours = load_safetensors_readonly(str(path)) + theirs = load_file(str(path)) + assert set(ours) == set(theirs) == set(tensors) + for name in tensors: + assert ours[name].dtype == theirs[name].dtype + assert ours[name].shape == theirs[name].shape + assert torch.equal(ours[name], theirs[name]) + assert safetensors_keys(str(path)) == list(theirs) + + +@pytest.mark.skipif(sys.platform != "linux", reason="/proc/self/maps") +def test_readonly_mapping_has_no_write_permission(tmp_path): + path, _ = _write(tmp_path) + tensor = dict(iter_safetensors_readonly(str(path)))["a.weight"] + ptr = tensor.data_ptr() + perms = None + for line in open("/proc/self/maps"): + fields = line.split() + low, high = (int(x, 16) for x in fields[0].split("-")) + if low <= ptr < high: + perms = fields[1] + break + assert perms is not None and perms.startswith("r--"), perms + + +def test_mmap_reader_maps_read_only_where_host_copies_are_redundant( + tmp_path, monkeypatch +): + from sglang.multimodal_gen.runtime.loader.weight_readers import safetensors_mmap + + path, tensors = _write(tmp_path) + monkeypatch.setattr(safetensors_mmap, "host_copies_are_redundant", lambda: True) + reader = safetensors_mmap.SafetensorsMmapReader() + got = dict( + reader.iter_weights( + [str(path)], + device="cpu", + to_cpu=True, + key_filter=lambda name: name != "b.bias", + show_progress=False, + ) + ) + assert set(got) == set(tensors) - {"b.bias"} + assert torch.equal(got["a.weight"], tensors["a.weight"]) + if sys.platform == "linux": + ptr = got["a.weight"].data_ptr() + perms = next( + line.split()[1] + for line in open("/proc/self/maps") + if int(line.split()[0].split("-")[0], 16) + <= ptr + < int(line.split()[0].split("-")[1], 16) + ) + assert perms.startswith("r--"), perms diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index f385875b7..5821a2caf 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py @@ -314,7 +314,7 @@ class TestDirectGPUVAEState(unittest.TestCase): "optimize_vae", side_effect=lambda vae: vae, ), - patch.object(vae_loader, "safetensors_load_file") as legacy_load, + patch("safetensors.torch.load_file") as legacy_load, ): safetensors_save_file( {"proj.weight": expected_weight, "scale": expected_scale},