From 0be2a209ac00c40786ef562cbc2418a9832e0dea Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 22 Aug 2026 09:37:31 +0800 Subject: [PATCH] [diffusion] refactor: hand out pinned host memory per layer (#35867) Co-authored-by: Claude Opus 5 --- .../memory_managers/layerwise_offload.py | 227 ++++++++++++----- .../test/unit/test_layerwise_offload.py | 238 ++++++++++++++++-- 2 files changed, 395 insertions(+), 70 deletions(-) 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 f0d5809ad..161a56e30 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 @@ -209,6 +209,8 @@ class LayerwiseOffloadManager: resident_layers: int = 0, initialize: bool = True, residency_policy: str = RESIDENCY_POLICY_LEADING, + pin_budget: HostPinBudget | None = None, + pin_component_name: str = "layerwise offload", ) -> None: self.model = model self.layers_attr_str = layers_attr_str @@ -217,6 +219,13 @@ class LayerwiseOffloadManager: # mps shares physical memory with the CPU and has no pinned host memory # or CUDA-style copy streams self.pin_cpu_memory = bool(pin_cpu_memory and not self._synchronous_mps) + # asked per layer rather than for the whole component; see + # _plan_pinned_layers + # A missing budget is not a licence to ignore host memory: without one + # every layer looked affordable and the copies-do-not-fit check below + # was never reached. A private budget reads the same host limit. + self._pin_budget = pin_budget if pin_budget is not None else HostPinBudget() + self._pin_component_name = pin_component_name # an explicit MPS zero avoids staging the next layer alongside the # active one; MPS has no transfer overlap to recover from that cost self.prefetch_size = ( @@ -276,7 +285,7 @@ class LayerwiseOffloadManager: # 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. + # not fit; see _plan_layer_hosting. 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 @@ -390,47 +399,127 @@ 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. + def _layer_byte_totals( + self, layer_groups: Dict + ) -> Tuple[Dict[int, int], Dict[int, int]]: + """Per layer: (all weight bytes, the subset that are checkpoint views).""" + totals: Dict[int, int] = {} + mapped: Dict[int, int] = {} + for layer_idx, dtype_to_params in layer_groups.items(): + total = 0 + from_mapping = 0 + for weights in dtype_to_params.values(): + for _, weight in weights: + tensor = self._to_local_tensor(weight) + nbytes = tensor.untyped_storage().nbytes() + total += nbytes + if self._mapped_regions.holds(tensor): + from_mapping += nbytes + totals[layer_idx] = total + mapped[layer_idx] = from_mapping + return totals, mapped - 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. + def _plan_layer_hosting(self, layer_groups: Dict) -> Dict[int, str]: + """Where each layer's weights live on the host: pinned, pageable or mapped. - 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. + Pinning is what lets the copy stream run ahead of compute; a pageable + or mapped source transfers synchronously however it is requested. + + The budget used to be asked for the whole component at once, so a DiT + larger than the whole spendable budget pinned nothing at all. Asking + per layer spends what there is. + + A layer that misses the budget falls back the way it always did, to a + pageable copy, and only stays on its mapping when those copies do not + fit either. The order matters: a pageable copy transfers synchronously, + since the driver stages it through its own pinned buffer, so it buys + none of the overlap -- but it is guaranteed resident, where a mapping can + be dropped and re-read from disk. + + Which layers get pinned matters only through how often each is read. + A streamed layer is transferred once per denoise step; a resident one is + transferred once per stage, so a pin on it is worth about 1/steps of the + same pin on a streamed layer. Streamed layers therefore get the budget + first, in streamed order, which is also deterministic. Unpinning a + resident layer costs one possibly-faulting arming copy per request and + buys a whole layer's worth of per-step overlap. """ - 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 + totals, mapped = self._layer_byte_totals(layer_groups) + pinned_bytes = 0 + hosting: Dict[int, str] = {} + pin_order: List[int] = [] + spendable = self._pin_budget.spendable_bytes if self._pin_budget else 0 + streamed = [idx for idx in self._streamed_order if idx in totals] + resident = [idx for idx in sorted(totals) if idx not in set(streamed)] + for layer_idx in streamed + resident: + layer_bytes = totals[layer_idx] + if self.pin_cpu_memory and pinned_bytes + layer_bytes <= spendable: + hosting[layer_idx] = "pinned" + pinned_bytes += layer_bytes + pin_order.append(layer_idx) + else: + hosting[layer_idx] = "pageable" + + def anonymous_new_bytes() -> int: + # What this plan adds, net, to anonymous memory. A store buffer + # that replaces an anonymous original -- the non-view share, such + # as a fused qkv -- is a wash: the original is freed when its + # parameter is rebound. The net cost of hosting a layer off its + # mapping is therefore only the checkpoint-view share it copies in, + # and a layer left on the mapping adds nothing. + return sum( + mapped[idx] for idx, where in hosting.items() if where != "mapped" + ) + + unpinned = [idx for idx, where in hosting.items() if where != "pinned"] + # The pins are booked but not yet allocated, so what the plan adds has + # to be weighed as one sum against the live reading. Asking about any + # one tier alone counts the same free bytes twice, and the error only + # ever says "fits". + if unpinned and host_copies_would_not_fit(anonymous_new_bytes()): + for layer_idx in unpinned: + if mapped[layer_idx]: + hosting[layer_idx] = "mapped" + # If the pins alone still do not fit, pins are what there is to + # give back. The tail of the pin order holds the least valuable + # ones, so they go first. + while pin_order and host_copies_would_not_fit(anonymous_new_bytes()): + layer_idx = pin_order.pop() + hosting[layer_idx] = "mapped" if mapped[layer_idx] else "pageable" + pinned_bytes -= totals[layer_idx] + if host_copies_would_not_fit(anonymous_new_bytes()): + logger.warning( + "Layerwise offload: %s adds %.2f GiB of anonymous host " + "memory that no mapping can absorb, and %.2f GiB is " + "available. Expect the host to be the limit.", + self._pin_component_name, + anonymous_new_bytes() / 1024**3, + host_memory_available_bytes() / 1024**3, + ) + if pinned_bytes and self._pin_budget is not None: + self._pin_budget.request( + component_name=self._pin_component_name, weight_bytes=pinned_bytes + ) + + if unpinned: + counts = {where: 0 for where in ("pinned", "pageable", "mapped")} + for where in hosting.values(): + counts[where] += 1 + logger.info( + "Layerwise offload: %s pins %d of %d layers (%.2f GiB of %.2f GiB " + "spendable). Of the rest, %d are copied into pageable host memory " + "and %d stay on the checkpoint mapping. Pinning every layer would " + "need %.2f GiB.", + self._pin_component_name, + counts["pinned"], + len(totals), + pinned_bytes / 1024**3, + spendable / 1024**3, + counts["pageable"], + counts["mapped"], + sum(totals.values()) / 1024**3, + ) + return hosting def _initialize_layer_weights(self) -> None: self._named_parameters = dict(self.model.named_parameters()) @@ -451,7 +540,7 @@ class LayerwiseOffloadManager: local_tensor.dtype, [] ).append((name, tensor)) - keep_mapping = self._keep_weights_on_their_mapping(layer_groups) + layer_hosting = self._plan_layer_hosting(layer_groups) # 2. concat and offload (in pinned memory) for layer_idx, dtype_to_params in layer_groups.items(): @@ -460,11 +549,14 @@ class LayerwiseOffloadManager: self._mapped_cpu_weights[layer_idx] = {} self._weight_metadata[layer_idx] = {} + hosting = layer_hosting.get(layer_idx, "pinned") + pin_this_layer = hosting == "pinned" + 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): + if hosting == "mapped" 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. @@ -499,7 +591,7 @@ class LayerwiseOffloadManager: size=local_weight.shape, stride=local_weight.stride(), dtype=dtype, - pin_memory=self.pin_cpu_memory, + pin_memory=pin_this_layer, ) cpu_tensor.copy_(local_weight) self._strided_cpu_weights[layer_idx][name] = cpu_tensor @@ -531,7 +623,7 @@ class LayerwiseOffloadManager: # create concatenated CPU buffer (in pinned memory) cpu_buffer = torch.empty( - total_numel, dtype=dtype, pin_memory=self.pin_cpu_memory + total_numel, dtype=dtype, pin_memory=pin_this_layer ) # offload weights to the buffer @@ -1239,20 +1331,20 @@ class LayerwiseOffloadableModuleMixin: # Pinning these weights is what lets the copy stream run ahead of # compute, but pinned pages are the ones the kernel cannot reclaim, - # so a component only gets them while the budget lasts. - pin_cpu_memory = server_args.pin_cpu_memory - if pin_cpu_memory and pin_budget is not None: - pin_cpu_memory = pin_budget.request( - component_name=f"{component_name or type(self).__name__}.{layer_name}", - weight_bytes=module_weight_bytes(module_list), - ) + # so they are handed out only while the budget lasts. The budget goes + # to the manager rather than being spent here, because it is asked + # per layer: a component too large to pin whole can still pin part of + # itself. See _plan_layer_hosting. + pin_component_name = f"{component_name or type(self).__name__}.{layer_name}" manager = LayerwiseOffloadManager( model=self, layers_attr_str=layer_name, num_layers=num_layers, enabled=True, - pin_cpu_memory=pin_cpu_memory, + pin_cpu_memory=server_args.pin_cpu_memory, + pin_budget=pin_budget, + pin_component_name=pin_component_name, prefetch_size=prefetch_size, resident_layers=resident_layers, initialize=False, @@ -1560,15 +1652,38 @@ def configure_layerwise_offload_modules( ) def _default_num_inference_steps() -> int: - from sglang.multimodal_gen.registry import get_pipeline_config_classes + from sglang.multimodal_gen.registry import ( + get_model_info, + get_pipeline_config_classes, + ) + sampling_cls = None pipeline_class_name = server_args.pipeline_class_name - if not pipeline_class_name: + if pipeline_class_name: + config_classes = get_pipeline_config_classes(pipeline_class_name) + if config_classes is not None: + sampling_cls = config_classes[1] + else: + # The override is normally unset. Resolve the pipeline the way + # build_pipeline does -- a cache hit by now -- because falling + # back to 1 here turns the benefit ranking into a bare-bytes + # ranking, and a once-per-request encoder can then outrank the + # stepped DiT for the pin budget. + model_path = getattr(server_args, "model_path", None) + if model_path: + model_info = get_model_info( + model_path, + backend=getattr(server_args, "backend", None), + model_id=getattr(server_args, "model_id", None), + ) + if model_info is not None: + sampling_cls = model_info.sampling_param_cls + if sampling_cls is None: return 1 - config_classes = get_pipeline_config_classes(pipeline_class_name) - if config_classes is None: + steps = getattr(sampling_cls(), "num_inference_steps", None) + if not steps: return 1 - return max(1, int(config_classes[1]().num_inference_steps)) + return max(1, int(steps)) default_steps = _default_num_inference_steps() 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 aaf4e69b3..2fb394904 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -429,6 +429,43 @@ def test_layerwise_pipeline_selection_uses_dit_group(monkeypatch): assert is_layerwise_offloaded_module(layerwise_module) +def test_pin_budget_ranks_by_steps_resolved_from_model_index(monkeypatch): + monkeypatch.setattr( + layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule + ) + monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu") + + import sglang.multimodal_gen.registry as registry_mod + + class _Sampling: + num_inference_steps = 50 + + monkeypatch.setattr( + registry_mod, + "get_model_info", + lambda *args, **kwargs: SimpleNamespace(sampling_param_cls=_Sampling), + ) + + # Same byte size on purpose: with the steps resolved, the stepped DiT + # outranks the encoder; with the silent steps=1 fallback the ranking + # ties and stable sort keeps the encoder first. + text_encoder = _NestedEncoderDummyModel() + transformer = _NestedDummyModel() + modules = {"text_encoder": text_encoder, "transformer": transformer} + + configured = configure_layerwise_offload_modules( + modules, + _server_args(model_path="/models/minimax-h3"), + component_names=["text_encoder", "transformer"], + ) + + assert configured == ["transformer", "text_encoder"], ( + "the pipeline is resolved from model_index when no override is set, " + "so the stepped DiT must claim the pin budget before the " + "once-per-request encoder" + ) + + def test_layerwise_configuration_filters_by_component_name(monkeypatch): monkeypatch.setattr( layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule @@ -1243,30 +1280,76 @@ class _FileBackedBlock(torch.nn.Module): class _FileBackedModel(torch.nn.Module): + def __init__(self, path: pathlib.Path, num_blocks: int = 1) -> None: + super().__init__() + self.blocks = torch.nn.ModuleList( + [ + _FileBackedBlock(path.with_name(f"{path.name}.{i}")) + for i in range(num_blocks) + ] + ) + + +# one _FileBackedBlock weight: 64 float32 +_BLOCK_BYTES = 64 * 4 + + +class _MixedBlock(torch.nn.Module): + """A block that is half checkpoint view, half anonymous memory -- the shape + of a layer whose qkv was fused at load while the rest stayed mapped.""" + def __init__(self, path: pathlib.Path) -> None: super().__init__() - self.blocks = torch.nn.ModuleList([_FileBackedBlock(path)]) + 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) + self.fused = torch.nn.Parameter( + torch.zeros(8, 8, dtype=torch.float32), requires_grad=False + ) -def _mapped_manager(tmp_path, monkeypatch, *, available_gib): +class _MixedModel(torch.nn.Module): + def __init__(self, path: pathlib.Path, num_blocks: int) -> None: + super().__init__() + self.blocks = torch.nn.ModuleList( + [_MixedBlock(path.with_name(f"{path.name}.{i}")) for i in range(num_blocks)] + ) + + +# one _MixedBlock: 64 float32 mapped + 64 float32 anonymous +_MIXED_BLOCK_BYTES = 2 * 64 * 4 + + +def _mapped_manager( + tmp_path, + monkeypatch, + *, + available_gib=None, + available_bytes=None, + num_blocks=1, + pin_budget_bytes=None, +): 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") + if available_bytes is None: + available_bytes = int(available_gib * 1024**3) + monkeypatch.setattr( + host_memory_budget, "host_memory_available_bytes", lambda: available_bytes + ) + model = _FileBackedModel(tmp_path / "weights.bin", num_blocks=num_blocks) return LayerwiseOffloadManager( model=model, layers_attr_str="blocks", - num_layers=1, + num_layers=num_blocks, enabled=True, pin_cpu_memory=True, + pin_budget=( + host_memory_budget.HostPinBudget(available_bytes=pin_budget_bytes) + if pin_budget_bytes is not None + else None + ), prefetch_size=1, ) @@ -1275,7 +1358,9 @@ def test_weights_stay_on_the_mapping_when_copies_do_not_fit(tmp_path, monkeypatc 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) + manager = _mapped_manager( + tmp_path, monkeypatch, available_gib=0.001, pin_budget_bytes=0 + ) 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) @@ -1292,7 +1377,9 @@ def test_weights_are_copied_when_they_fit(tmp_path, monkeypatch): 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) + manager = _mapped_manager( + tmp_path, monkeypatch, available_gib=0.001, pin_budget_bytes=0 + ) before = manager._mapped_cpu_weights[0]["blocks.0.weight"].clone() manager._gpu_layers.add(0) manager.sync_layer_to_cpu(0) @@ -1331,10 +1418,133 @@ def test_the_mapped_store_survives_the_placeholder(tmp_path, monkeypatch): ), "the byte counter and the store must describe the same weights" +def test_only_the_layers_the_budget_covers_are_pinned(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + # The pin budget covers two of the four layers. Available host memory sits + # above the copy reserve by less than all four layers but more than the + # two pins, so the rest is demoted to the mapping and the pins stand. + budget = 2 * 1024**3 + 2 * _BLOCK_BYTES + manager = _mapped_manager( + tmp_path, + monkeypatch, + available_bytes=4 * 1024**3 + 3 * _BLOCK_BYTES, + pin_budget_bytes=budget, + num_blocks=4, + ) + pinned = {i for i in range(4) if manager._consolidated_cpu_weights.get(i)} + mapped = {i for i in range(4) if manager._mapped_cpu_weights.get(i)} + assert pinned == {0, 1}, "the layers the budget covers, taken in index order" + assert mapped == {2, 3} + assert not (pinned & mapped), "a layer is in one store or the other" + + +def test_pins_are_given_back_when_they_do_not_fit_the_host(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + # Four fully mapped layers. The budget covers two pins, but each pin copies + # its whole layer off the mapping and the host only has room for one plus + # the reserve: the plan must give the second pin back rather than allocate + # more than the machine has. + available = 4 * 1024**3 + int(1.5 * _BLOCK_BYTES) + manager = _mapped_manager( + tmp_path, + monkeypatch, + available_bytes=available, + pin_budget_bytes=2 * 1024**3 + 2 * _BLOCK_BYTES, + num_blocks=4, + ) + on_mapping = {i for i in range(4) if manager._mapped_cpu_weights.get(i)} + assert on_mapping == {1, 2, 3}, ( + "two pins add two layers of anonymous memory against room for 1.5, so " + "the least valuable pin (layer 1) is given back" + ) + pinned = {i for i in range(4) if manager._consolidated_cpu_weights.get(i)} + assert pinned == {0} + + +def test_replacing_an_anonymous_original_is_not_charged_as_new(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + # Half of each layer is a fused anonymous tensor whose store buffer + # replaces it -- a wash. Only the mapped half is a net addition, so a + # host with room for the mapped halves alone must still pin every layer. + block = _MIXED_BLOCK_BYTES + available = 4 * 1024**3 + int(2.5 * block) + monkeypatch.setattr( + layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule + ) + monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu") + monkeypatch.setattr( + host_memory_budget, "host_memory_available_bytes", lambda: available + ) + model = _MixedModel(tmp_path / "weights.bin", num_blocks=4) + manager = LayerwiseOffloadManager( + model=model, + layers_attr_str="blocks", + num_layers=4, + enabled=True, + pin_cpu_memory=True, + pin_budget=host_memory_budget.HostPinBudget( + available_bytes=2 * 1024**3 + 3 * block + ), + prefetch_size=1, + ) + pinned = {i for i in range(4) if 0 in manager._consolidated_cpu_weights} + on_mapping = {i for i in range(4) if manager._mapped_cpu_weights.get(i)} + assert not on_mapping, ( + "the net addition is three pinned mapped-halves plus one pageable " + "mapped-half (2 blocks) against room for 2.5 -- charging the replaced " + "anonymous originals as new would demote, then strip every pin" + ) + assert all(manager._consolidated_cpu_weights.get(i) for i in range(4)) + + +def test_every_layer_is_pinned_when_the_budget_covers_them(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, + pin_budget_bytes=64 * 1024**3, + num_blocks=4, + ) + assert not any(manager._mapped_cpu_weights.get(i) for i in range(4)) + assert all(manager._consolidated_cpu_weights.get(i) for i in range(4)) + + +def test_no_layer_is_pinned_when_the_budget_is_spent(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_bytes=0, pin_budget_bytes=0, num_blocks=4 + ) + assert all(manager._mapped_cpu_weights.get(i) for i in range(4)) + + +def test_an_unpinnable_layer_is_still_copied_when_the_copies_fit(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + # no pinned budget at all, but plenty of host memory: a pageable copy is + # guaranteed resident where a mapping can be dropped and re-read + manager = _mapped_manager( + tmp_path, + monkeypatch, + available_gib=64, + pin_budget_bytes=0, + num_blocks=4, + ) + assert not any(manager._mapped_cpu_weights.get(i) for i in range(4)) + assert all(manager._consolidated_cpu_weights.get(i) for i in range(4)) + + 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) + manager = _mapped_manager( + tmp_path, monkeypatch, available_gib=0.001, pin_budget_bytes=0 + ) names = {name for name, _ in manager.iter_cpu_weights()} assert "blocks.0.weight" in names