[diffusion] refactor: hand out pinned host memory per layer (#35867)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

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:
+    def _layer_byte_totals(
-        """Whether to leave file-backed weights on their mapping.
+        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
+    def _plan_layer_hosting(self, layer_groups: Dict) -> Dict[int, str]:
-        copy stream run ahead of compute -- worth 1.03 s against 1.90 s per step
+        """Where each layer's weights live on the host: pinned, pageable or mapped.
-        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
+        Pinning is what lets the copy stream run ahead of compute; a pageable
-        weights, of which 50.53 GiB arrive as views into the checkpoint; on a
+        or mapped source transfers synchronously however it is requested.
-        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
+        The budget used to be asked for the whole component at once, so a DiT
-        nearly as fast as pinned -- 12.38 GB/s against 13.39 measured -- and what
+        larger than the whole spendable budget pinned nothing at all. Asking
-        is lost is the overlap, not the bandwidth.
+        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:
+        totals, mapped = self._layer_byte_totals(layer_groups)
-            return False
+        pinned_bytes = 0
-        mapped_bytes = sum(
+        hosting: Dict[int, str] = {}
-            tensor.untyped_storage().nbytes()
+        pin_order: List[int] = []
-            for dtype_to_params in layer_groups.values()
+        spendable = self._pin_budget.spendable_bytes if self._pin_budget else 0
-            for weights in dtype_to_params.values()
+        streamed = [idx for idx in self._streamed_order if idx in totals]
-            for _, weight in weights
+        resident = [idx for idx in sorted(totals) if idx not in set(streamed)]
-            for tensor in (self._to_local_tensor(weight),)
+        for layer_idx in streamed + resident:
-            if self._mapped_regions.holds(tensor)
+            layer_bytes = totals[layer_idx]
-        )
+            if self.pin_cpu_memory and pinned_bytes + layer_bytes <= spendable:
-        if mapped_bytes <= 0:
+                hosting[layer_idx] = "pinned"
-            return False
+                pinned_bytes += layer_bytes
-        # The copies land in host memory on top of the page cache already holding
+                pin_order.append(layer_idx)
-        # these bytes, so the room needed is the copy itself plus a reserve.
+            else:
-        if not host_copies_would_not_fit(mapped_bytes):
+                hosting[layer_idx] = "pageable"
-            return False
+
-        available = host_memory_available_bytes()
+        def anonymous_new_bytes() -> int:
-        logger.info(
+            # What this plan adds, net, to anonymous memory. A store buffer
-            "Layerwise offload: leaving %.2f GiB of weights on the checkpoint "
+            # that replaces an anonymous original -- the non-view share, such
-            "mapping -- copying them into host memory needs more than the "
+            # as a fused qkv -- is a wash: the original is freed when its
-            "%.2f GiB available. The page cache decides what stays resident, so "
+            # parameter is rebound. The net cost of hosting a layer off its
-            "reads may come from disk; the copies cannot be pinned, so they run "
+            # mapping is therefore only the checkpoint-view share it copies in,
-            "on the compute stream instead of ahead of it.",
+            # and a layer left on the mapping adds nothing.
-            mapped_bytes / 1024**3,
+            return sum(
-            available / 1024**3,
+                mapped[idx] for idx, where in hosting.items() if where != "mapped"
-        )
+            )
-        return True
+
+        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.
+            # so they are handed out only while the budget lasts. The budget goes
-            pin_cpu_memory = server_args.pin_cpu_memory
+            # to the manager rather than being spent here, because it is asked
-            if pin_cpu_memory and pin_budget is not None:
+            # per layer: a component too large to pin whole can still pin part of
-                pin_cpu_memory = pin_budget.request(
+            # itself. See _plan_layer_hosting.
-                    component_name=f"{component_name or type(self).__name__}.{layer_name}",
+            pin_component_name = f"{component_name or type(self).__name__}.{layer_name}"
-                    weight_bytes=module_weight_bytes(module_list),
-                )
 
             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)
+        steps = getattr(sampling_cls(), "num_inference_steps", None)
-        if config_classes is 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()
 

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
+    if available_bytes is None:
-    # whether the copies fit, and layerwise_offload's own is what the log reports.
+        available_bytes = int(available_gib * 1024**3)
-    available_bytes = int(available_gib * 1024**3)
+    monkeypatch.setattr(
-    for module in (host_memory_budget, layerwise_offload_mod):
+        host_memory_budget, "host_memory_available_bytes", lambda: available_bytes
-        monkeypatch.setattr(
+    )
-            module, "host_memory_available_bytes", lambda: available_bytes
+    model = _FileBackedModel(tmp_path / "weights.bin", num_blocks=num_blocks)
-        )
-    model = _FileBackedModel(tmp_path / "weights.bin")
     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