minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Mick
2026-08-22 09:37:31 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent b26695a26e
commit 0be2a209ac
2 changed files with 395 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py
+171 -56
View File
@@ -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()
python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py
+224 -14
View File
@@ -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