From 6127d1daeee33ec758a737bca0a1ddd2d17f02ca Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 21 Aug 2026 10:54:28 +0800 Subject: [PATCH] [diffusion] feat: allow offloaded weights stay on the checkpoint mapping (#35701) --- .../component_loaders/text_encoder_loader.py | 48 +++++++- .../runtime/loader/fsdp_load.py | 18 +-- .../multimodal_gen/runtime/loader/utils.py | 48 ++++++-- .../memory_managers/host_memory_budget.py | 17 +++ .../memory_managers/layerwise_offload.py | 104 +++++++++++++++++- .../runtime/models/dits/minimax_h3.py | 2 +- .../models/encoders/minimax_h3_qwen3vl.py | 2 +- .../test/unit/test_host_memory_budget.py | 45 ++++++++ .../test/unit/test_layerwise_offload.py | 80 ++++++++++++++ 9 files changed, 338 insertions(+), 26 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 69ee4800f..dd870b338 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -39,6 +39,10 @@ from sglang.multimodal_gen.runtime.loader.weight_utils import ( pt_weights_iterator, safetensors_weights_iterator, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + host_copies_would_not_fit, + host_memory_available_bytes, +) from sglang.multimodal_gen.runtime.models.encoders.base import ( EncoderTensorParallelMixin, TextEncoder, @@ -205,6 +209,34 @@ def _process_quantized_encoder_weights( return processed_layers +def _checkpoint_bytes(model_path: str) -> int: + """On-disk size of a checkpoint, readable before any weight of it is.""" + total = 0 + for path in glob.glob( + os.path.join(str(model_path), "**", "*.safetensors"), recursive=True + ): + try: + total += os.path.getsize(path) + except OSError: + continue + return total + + +def _keep_this_checkpoint_mapped(model_path: str) -> bool: + """Whether this encoder's weights should stay on their file mapping.""" + checkpoint_bytes = _checkpoint_bytes(model_path) + if not host_copies_would_not_fit(checkpoint_bytes): + return False + logger.info( + "Text encoder checkpoint is %.2f GiB against %.2f GiB of host memory, " + "so its compatible weights stay on the checkpoint mapping instead of " + "being copied in.", + checkpoint_bytes / 1024**3, + host_memory_available_bytes() / 1024**3, + ) + return True + + class TextEncoderLoader(ComponentLoader): """Loader for text encoders.""" @@ -613,11 +645,17 @@ class TextEncoderLoader(ComponentLoader): ) model.bind_encoder_tp_group(encoder_tp_group) - if current_platform.is_mps() and component_starts_on_cpu: - # the h3 encoder is layered immediately after this loader returns - # compatible CPU safetensors stay mapped instead of copying the - # full Qwen checkpoint into unified memory - model._mps_zero_copy_weight_loading = True + if component_starts_on_cpu and ( + current_platform.is_mps() or _keep_this_checkpoint_mapped(model_path) + ): + # The encoder is layered immediately after this loader returns, + # so compatible CPU safetensors can stay mapped instead of being + # copied. On MPS that is always the right call -- the memory is + # unified. On any host it becomes the only call once the + # checkpoint is larger than host memory, because the copy is + # what does not fit: H3's encoder is 62.13 GiB against a 32 GiB + # target. + model._keep_checkpoint_mapping = True weights_to_load = {name for name, _ in model.named_parameters()} loaded_weights = model.load_weights( diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 3f7503c37..ca53f18cc 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -299,15 +299,15 @@ def maybe_load_fsdp_model( logger.info("Disabling FSDP for MPS platform as it's not compatible") weight_load_plan = weight_load_plan or WeightLoadPlan(checkpoint_load_device=device) - mps_zero_copy_weight_loading = bool( + keep_checkpoint_mapping = bool( current_platform.is_mps() and weight_load_plan.mps_layerwise_cpu_staging and weight_load_plan.checkpoint_load_device.type == "cpu" ) - if mps_zero_copy_weight_loading: - # layerwise offload replaces block parameters with mps placeholders after + if keep_checkpoint_mapping: + # layerwise offload replaces block parameters with placeholders after # load, so compatible checkpoint tensors stay file-backed on CPU - model._mps_zero_copy_weight_loading = True + model._keep_checkpoint_mapping = True defer_cpu_placement = bool( component_starts_on_cpu and weight_load_plan.defer_cpu_placement @@ -425,7 +425,7 @@ def maybe_load_fsdp_model( strict=strict, cpu_offload=load_on_cpu, param_names_mapping=param_names_mapping_fn, - mps_zero_copy_weight_loading=mps_zero_copy_weight_loading, + keep_checkpoint_mapping=keep_checkpoint_mapping, preconverted_state_dict=preconverted_state_dict, ) if bnb_quant_states: @@ -549,7 +549,7 @@ def load_model_from_full_model_state_dict( strict: bool = False, cpu_offload: bool = False, param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None, - mps_zero_copy_weight_loading: bool = False, + keep_checkpoint_mapping: bool = False, preconverted_state_dict: ( tuple[ dict[ @@ -574,7 +574,7 @@ def load_model_from_full_model_state_dict( strict (bool): flag to check if to load the model in strict mode cpu_offload (bool): flag to check if FSDP offload is enabled param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name - mps_zero_copy_weight_loading (bool): retain compatible CPU checkpoint tensors for MPS layerwise offload + keep_checkpoint_mapping (bool): retain compatible CPU checkpoint tensors instead of copying them Returns: ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: * **missing_keys** is a list of str containing the missing keys @@ -717,9 +717,9 @@ def load_model_from_full_model_state_dict( else None ) use_checkpoint_tensor_directly = bool( - mps_zero_copy_weight_loading + keep_checkpoint_mapping and actual_param is not None - and not getattr(actual_param, "mps_zero_copy_unsafe", False) + and not getattr(actual_param, "checkpoint_mapping_unsafe", False) and tuple(meta_sharded_param.shape) == tuple(full_tensor.shape) and full_tensor.device.type == "cpu" and full_tensor.dtype == target_dtype diff --git a/python/sglang/multimodal_gen/runtime/loader/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index a70ed6974..704a81e39 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -357,6 +357,44 @@ def _read_process_mappings() -> tuple[list[int], list[int], list[bool]] | None: return [r[0] for r in rows], [r[1] for r in rows], [r[2] for r in rows] +class MappedRegions: + """Answers whether a tensor's bytes live in a file mapping. + + Built once and reused. The lookup table comes from /proc/self/maps, so + rebuilding it per tensor would be quadratic over a checkpoint's worth of + weights -- H3's DiT alone has tens of thousands. + + A snapshot, not a live view: mappings created after construction are + unknown to it. Callers that need to classify freshly loaded weights should + build one after loading, which is when the mappings exist. + """ + + def __init__(self) -> None: + self._maps = _read_process_mappings() + + @property + def available(self) -> bool: + """False where /proc is absent, in which case nothing is classified.""" + return self._maps is not None + + def holds_pointer(self, pointer: int) -> bool: + if self._maps is None or pointer == 0: + return False + starts, ends, backed = self._maps + index = bisect.bisect_right(starts, pointer) - 1 + if index < 0 or pointer >= ends[index]: + return False + return backed[index] + + def holds(self, tensor: torch.Tensor) -> bool: + if tensor.device.type != "cpu": + return False + try: + return self.holds_pointer(tensor.untyped_storage().data_ptr()) + except Exception: + return False + + def component_residency_bytes(module) -> Dict[str, int]: """Where a component's weights actually sit, in bytes. @@ -380,16 +418,10 @@ def component_residency_bytes(module) -> Dict[str, int]: totals = {"vram": 0, "host_pinned": 0, "host_mapped": 0, "host": 0} seen: set[int] = set() - mappings = _read_process_mappings() + regions = MappedRegions() def is_file_backed(pointer: int) -> bool: - if mappings is None: - return False - starts, ends, backed = mappings - index = bisect.bisect_right(starts, pointer) - 1 - if index < 0 or pointer >= ends[index]: - return False - return backed[index] + return regions.holds_pointer(pointer) def add(tensor: torch.Tensor) -> None: try: 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 a3f1d00f6..8b057cc29 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 @@ -41,6 +41,10 @@ _UNLIMITED_ABOVE = 1 << 62 HOST_RESERVE_FRACTION = 0.05 MIN_HOST_RESERVE_BYTES = 2 * GIB_BYTES +# Left free when weighing a checkpoint against host memory: activations, staging +# buffers and allocator slack are none of them in the weight total. +HOST_COPY_RESERVE_BYTES = 4 * GIB_BYTES + def _read_int(path: str) -> int | None: try: @@ -134,6 +138,19 @@ def host_memory_available_bytes() -> int: return min(available, max(0, limit - usage)) +def host_copies_would_not_fit(weight_bytes: int) -> bool: + """Whether copying `weight_bytes` into host memory would run the host out. + + The alternative to a copy is leaving the weights on their file mapping, + which the kernel may drop under pressure and re-read from disk. That is + slower per byte but bounded, so it is the right answer exactly when the + copies do not fit -- and the wrong one when they do. + """ + if weight_bytes <= 0: + return False + return weight_bytes >= host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES + + class HostPinBudget: """Hands out pinned-host-memory allowances until the headroom runs out. 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 e3f95ddf5..b7474873b 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 @@ -8,6 +8,7 @@ import torch from torch.distributed.tensor import DTensor from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType +from sglang.multimodal_gen.runtime.loader.utils import MappedRegions from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_RESIDENCY_GROUPS, LAYERWISE_OFFLOAD, @@ -16,6 +17,8 @@ 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_would_not_fit, + host_memory_available_bytes, module_weight_bytes, pin_benefit_bytes, ) @@ -270,6 +273,12 @@ class LayerwiseOffloadManager: # layer_idx -> {name: pinned_cpu_tensor_with_original_stride} # stores tensors whose original non-contiguous stride/layout must be preserved self._strided_cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {} + # layer_idx -> {name: tensor still viewing the checkpoint file} + # Weights left on their mapping rather than copied into host memory, so + # the page cache decides what stays resident. Used when the copies would + # not fit; see _keep_weights_on_their_mapping. + self._mapped_cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {} + self._mapped_bytes = 0 # mps keeps the original CPU tensor for each layer instead of building a # second flattened host copy self._mps_cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {} @@ -285,6 +294,9 @@ class LayerwiseOffloadManager: self._named_buffers: Dict[str, torch.Tensor] = {} self._offload_placeholders: Dict[torch.dtype, torch.Tensor] = {} self._has_dtensor_weights = False + # A snapshot of this process's mappings, taken now because the weights + # have just been loaded and their mappings exist. + self._mapped_regions = MappedRegions() # Store forward hooks for removal self._forward_hooks: List[Any] = [] @@ -378,6 +390,48 @@ class LayerwiseOffloadManager: self._finalize_initialization() + def _keep_weights_on_their_mapping(self, layer_groups: Dict) -> bool: + """Whether to leave file-backed weights on their mapping. + + Copying them into host memory buys pinning, and pinning is what lets the + copy stream run ahead of compute -- worth 1.03 s against 1.90 s per step + on a measured Wan2.1 run, so it is not given up lightly. + + It is given up when the copies do not fit. H3's DiT is 61.73 GiB of + weights, of which 50.53 GiB arrive as views into the checkpoint; on a + 32 GiB host the copy cannot be made at all, and the choice is between a + mapping and not running. Above that, page-cache residency makes the read + nearly as fast as pinned -- 12.38 GB/s against 13.39 measured -- and what + is lost is the overlap, not the bandwidth. + """ + if not self._mapped_regions.available: + return False + mapped_bytes = sum( + tensor.untyped_storage().nbytes() + for dtype_to_params in layer_groups.values() + for weights in dtype_to_params.values() + for _, weight in weights + for tensor in (self._to_local_tensor(weight),) + if self._mapped_regions.holds(tensor) + ) + if mapped_bytes <= 0: + return False + # The copies land in host memory on top of the page cache already holding + # these bytes, so the room needed is the copy itself plus a reserve. + if not host_copies_would_not_fit(mapped_bytes): + return False + available = host_memory_available_bytes() + logger.info( + "Layerwise offload: leaving %.2f GiB of weights on the checkpoint " + "mapping -- copying them into host memory needs more than the " + "%.2f GiB available. The page cache decides what stays resident, so " + "reads may come from disk; the copies cannot be pinned, so they run " + "on the compute stream instead of ahead of it.", + mapped_bytes / 1024**3, + available / 1024**3, + ) + return True + def _initialize_layer_weights(self) -> None: self._named_parameters = dict(self.model.named_parameters()) self._named_buffers = dict(self.model.named_buffers()) @@ -397,16 +451,36 @@ class LayerwiseOffloadManager: local_tensor.dtype, [] ).append((name, tensor)) + keep_mapping = self._keep_weights_on_their_mapping(layer_groups) + # 2. concat and offload (in pinned memory) for layer_idx, dtype_to_params in layer_groups.items(): self._consolidated_cpu_weights[layer_idx] = {} self._strided_cpu_weights[layer_idx] = {} + self._mapped_cpu_weights[layer_idx] = {} self._weight_metadata[layer_idx] = {} for dtype, weights in dtype_to_params.items(): contiguous_weights: List[Tuple[str, torch.Tensor, torch.Tensor]] = [] for name, weight in weights: local_weight = self._to_local_tensor(weight) + if keep_mapping and self._mapped_regions.holds(local_weight): + # Already a view into the checkpoint. Copying it would + # add a second copy of bytes the page cache holds + # anyway, and that copy is what does not fit. + self._mapped_cpu_weights[layer_idx][name] = local_weight + self._weight_metadata[layer_idx][name] = { + "dtype": local_weight.dtype, + "shape": tuple(local_weight.shape), + "stride": local_weight.stride(), + "preserve_strides": False, + "mapped": True, + } + self._mapped_bytes += local_weight.untyped_storage().nbytes() + weight.data = self._get_shared_empty_tensor_for_target( + weight, local_weight.dtype + ) + continue if local_weight.is_contiguous(): contiguous_weights.append((name, weight, local_weight)) continue @@ -629,7 +703,9 @@ class LayerwiseOffloadManager: ) self._gpu_layers.add(layer_idx) return - if layer_idx not in self._consolidated_cpu_weights: + if layer_idx not in self._consolidated_cpu_weights and not ( + self._mapped_cpu_weights.get(layer_idx) + ): return if self.copy_stream is not None: self.copy_stream.wait_stream(torch.get_device_module().current_stream()) @@ -646,7 +722,9 @@ class LayerwiseOffloadManager: torch.no_grad(), stream_context, ): - for dtype, cpu_buffer in self._consolidated_cpu_weights[layer_idx].items(): + for dtype, cpu_buffer in self._consolidated_cpu_weights.get( + layer_idx, {} + ).items(): gpu_buffer = torch.empty( cpu_buffer.shape, dtype=dtype, device=self.device ) @@ -657,6 +735,18 @@ class LayerwiseOffloadManager: # so the recorded event covers both flat-buffer and stride-preserving copies. for name, meta in self._weight_metadata[layer_idx].items(): target = self.get_target_with_name(name) + if meta.get("mapped", False): + # Straight from the mapping. Not pinned, so this copy runs + # on the compute stream 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] + gpu_tensor = torch.empty( + meta["shape"], dtype=meta["dtype"], device=self.device + ) + gpu_tensor.copy_(cpu_tensor, non_blocking=False) + target.data = self._wrap_for_target(target, gpu_tensor) + continue + if meta.get("preserve_strides", False): # Recreate the original view layout instead of flatten+view. # ModelOpt FP8 relies on a transposed runtime weight layout, @@ -766,6 +856,12 @@ class LayerwiseOffloadManager: # Collect current GPU weights and write back to CPU buffer for name, meta in self._weight_metadata.get(layer_idx, {}).items(): + if meta.get("mapped", False): + # The store is the checkpoint file. Inference does not mutate + # weights, so there is nothing to write back; writing would + # copy-on-write the mapping into the anonymous memory this + # exists to avoid. + continue target = self.get_target_with_name(name) target_local = self._to_local_tensor(target) if meta.get("preserve_strides", False): @@ -903,6 +999,10 @@ class LayerwiseOffloadManager: for layer_idx in sorted(self._weight_metadata): for name, meta in self._weight_metadata[layer_idx].items(): + if meta.get("mapped", False): + yield name, self._mapped_cpu_weights[layer_idx][name] + continue + if meta.get("preserve_strides", False): # Some quantized weights rely on a non-contiguous layout. # Yield the strided tensor directly instead of rebuilding it diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 2907a4ea4..822a80ec8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -628,7 +628,7 @@ class MiniMaxH3Attention(nn.Module): weight = self.qkv_proj.weight # h3 checkpoints interleave each attention head's Q, K, and V rows # this parameter needs reordering before the native QKV projection - weight.mps_zero_copy_unsafe = True + weight.checkpoint_mapping_unsafe = True base_loader = weight.weight_loader def _reorder_checkpoint_weight(loaded_weight: torch.Tensor) -> torch.Tensor: diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py index a23835a3e..8f0b71724 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py @@ -200,7 +200,7 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder): weight_loader = getattr(param, "weight_loader", default_weight_loader) try: can_keep_checkpoint_tensor = bool( - getattr(self, "_mps_zero_copy_weight_loading", False) + getattr(self, "_keep_checkpoint_mapping", False) and weight_loader is default_weight_loader and param.device.type == "cpu" and loaded_weight.device.type == "cpu" diff --git a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py index 24182821b..3b8d99737 100644 --- a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py @@ -6,8 +6,10 @@ import torch.nn as nn from sglang.multimodal_gen.runtime.managers.memory_managers import host_memory_budget from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( GIB_BYTES, + HOST_COPY_RESERVE_BYTES, HostPinBudget, cgroup_memory_limit_bytes, + host_copies_would_not_fit, host_memory_available_bytes, module_weight_bytes, pin_benefit_bytes, @@ -98,6 +100,49 @@ class TestCgroupLimit: assert host_memory_available_bytes() == 12 * GIB_BYTES +def _host_has(monkeypatch, tmp_path, available_bytes): + _point_at(monkeypatch, tmp_path) + monkeypatch.setattr( + host_memory_budget.psutil, + "virtual_memory", + lambda: type("VM", (), {"available": available_bytes})(), + ) + + +class TestHostCopiesWouldNotFit: + def test_a_checkpoint_larger_than_the_host_does_not_fit( + self, monkeypatch, tmp_path + ): + # the measured H3 case: a 62 GiB encoder against a 32 GiB host + _host_has(monkeypatch, tmp_path, 32 * GIB_BYTES) + assert host_copies_would_not_fit(62 * GIB_BYTES) + + def test_a_checkpoint_well_under_the_host_fits(self, monkeypatch, tmp_path): + _host_has(monkeypatch, tmp_path, 128 * GIB_BYTES) + assert not host_copies_would_not_fit(20 * GIB_BYTES) + + def test_the_reserve_is_not_available_to_the_copy(self, monkeypatch, tmp_path): + # a copy that fits only by eating the reserve is treated as not fitting + _host_has(monkeypatch, tmp_path, 40 * GIB_BYTES) + spendable = 40 * GIB_BYTES - HOST_COPY_RESERVE_BYTES + assert host_copies_would_not_fit(spendable) + assert not host_copies_would_not_fit(spendable - 1) + + def test_nothing_to_copy_always_fits(self, monkeypatch, tmp_path): + _host_has(monkeypatch, tmp_path, 0) + assert not host_copies_would_not_fit(0) + + def test_the_cgroup_cap_decides_and_not_the_machine(self, monkeypatch, tmp_path): + # psutil sees the whole box; only the cap makes the copy impossible + _point_at(monkeypatch, tmp_path, v2=(32 * GIB_BYTES, 0)) + monkeypatch.setattr( + host_memory_budget.psutil, + "virtual_memory", + lambda: type("VM", (), {"available": 900 * GIB_BYTES})(), + ) + assert host_copies_would_not_fit(62 * GIB_BYTES) + + class TestNestedCgroup: def test_a_tighter_nested_cap_wins_over_the_root(self, monkeypatch, tmp_path): # a systemd scope with MemoryMax, or --cgroup-parent: planning against diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index d5f2a77bd..8f68fc35d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -1,3 +1,4 @@ +import pathlib from contextlib import nullcontext from types import SimpleNamespace @@ -15,6 +16,9 @@ from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( from sglang.multimodal_gen.runtime.managers.memory_managers import ( component_residency_strategies as component_residency_strategies_mod, ) +from sglang.multimodal_gen.runtime.managers.memory_managers import ( + host_memory_budget, +) from sglang.multimodal_gen.runtime.managers.memory_managers import ( layerwise_offload as layerwise_offload_mod, ) @@ -1228,6 +1232,82 @@ def test_strided_forward_leaves_exactly_the_resident_set(monkeypatch): assert len(manager._gpu_layers) <= len(resident) + manager.prefetch_size +class _FileBackedBlock(torch.nn.Module): + """A block whose weight is a view into a file, as a loaded checkpoint is.""" + + def __init__(self, path: pathlib.Path) -> None: + super().__init__() + path.write_bytes(b"\x00" * (64 * 4)) + mapped = torch.from_file(str(path), shared=True, size=64, dtype=torch.float32) + self.weight = torch.nn.Parameter(mapped.reshape(8, 8), requires_grad=False) + + +class _FileBackedModel(torch.nn.Module): + def __init__(self, path: pathlib.Path) -> None: + super().__init__() + self.blocks = torch.nn.ModuleList([_FileBackedBlock(path)]) + + +def _mapped_manager(tmp_path, monkeypatch, *, available_gib): + monkeypatch.setattr( + layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule + ) + monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu") + # Two bindings for one fact: the budget module's copy is what decides + # whether the copies fit, and layerwise_offload's own is what the log reports. + available_bytes = int(available_gib * 1024**3) + for module in (host_memory_budget, layerwise_offload_mod): + monkeypatch.setattr( + module, "host_memory_available_bytes", lambda: available_bytes + ) + model = _FileBackedModel(tmp_path / "weights.bin") + return LayerwiseOffloadManager( + model=model, + layers_attr_str="blocks", + num_layers=1, + enabled=True, + pin_cpu_memory=True, + prefetch_size=1, + ) + + +def test_weights_stay_on_the_mapping_when_copies_do_not_fit(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + # the reserve alone exceeds this, so no copy can be afforded + manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001) + assert manager._mapped_cpu_weights[0], "expected the weight to stay mapped" + assert manager._weight_metadata[0]["blocks.0.weight"]["mapped"] is True + assert not manager._consolidated_cpu_weights.get(0) + + +def test_weights_are_copied_when_they_fit(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + manager = _mapped_manager(tmp_path, monkeypatch, available_gib=64) + assert not manager._mapped_cpu_weights[0], "a copy was affordable" + assert manager._consolidated_cpu_weights[0] + + +def test_a_mapped_weight_is_not_written_back(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001) + before = manager._mapped_cpu_weights[0]["blocks.0.weight"].clone() + manager._gpu_layers.add(0) + manager.sync_layer_to_cpu(0) + after = manager._mapped_cpu_weights[0]["blocks.0.weight"] + assert torch.equal(before, after), "writeback must not touch the checkpoint" + + +def test_mapped_weights_are_visible_to_checksums(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001) + names = {name for name, _ in manager.iter_cpu_weights()} + assert "blocks.0.weight" in names + + def test_layerwise_tuning_defaults_match_the_group(): """No per-component entry: the DiT group keeps its knobs, auxiliaries do not.""" args = _server_args(