From ca0d8271da4efaa256ac1af668814620b7995db7 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Wed, 19 Aug 2026 18:33:50 -0600 Subject: [PATCH 1/2] pgw#1515: coarse residency cells + the hold-while-hot schedule (checkpoint) The packing and scheduling halves of phase 3's residency work, both PLANNER-side so they sit above the mechanism seam and serve software streaming and the varena facade alike. * Cells: pack_cells groups leaves into back/unback/stream units at a configurable target size, small leaves first. The plan reports the residual granularity tax against the leaf-granular baseline pgw#1507 measured (13.2%), so the elimination is a subtraction of two published numbers. * Schedule: plan_residency decides CALL_BOUNDARY vs PER_STEP from the assigned {VRAM, RAM} pair and the cell layout. hold_component() is the call-boundary swap and REFUSES when the plan did not admit it. * The default geometry is leaf-granular with no page granularity, so every pre-cell number is unchanged. --- src/gen_worker/models/stream_residency.py | 531 +++++++++++++++++-- tests/test_residency_cells.py | 596 ++++++++++++++++++++++ 2 files changed, 1097 insertions(+), 30 deletions(-) create mode 100644 tests/test_residency_cells.py diff --git a/src/gen_worker/models/stream_residency.py b/src/gen_worker/models/stream_residency.py index 0b4d39a6b..874fa859c 100644 --- a/src/gen_worker/models/stream_residency.py +++ b/src/gen_worker/models/stream_residency.py @@ -51,6 +51,34 @@ are forced resident (they are megabytes, and their cost is dwarfed by the launch overhead of streaming them). A MERGED adapter is just a different weight and streams like any other. + +**Cells and the schedule** (pgw#1515). Two things above the split, both +PLANNER-side and therefore shared by every mechanism the rung can sit on: + +* **Cells** (:class:`CellPolicy`, :class:`ResidencyCell`) — the unit of + back/unback/stream is a CELL, a packed group of leaves, not a leaf. A + mechanism that maps memory at a page granularity pays that granularity ONCE + PER REGION, so per-leaf regions charge the alignment remainder once per leaf: + pgw#1507 measured **248 MiB = 13.2 %** of sd1.5's weights lost that way + (1.830 GiB of weight in 2.072 GiB of span over 239 regions, 2 MiB VMM + granularity), and ~98 µs of page-table work per streamed leaf per forward on + top. Packing the small leaves first collapses both: the remainder is paid + once per CELL, and the plan REPORTS the residual so the elimination is a + measured claim and not an assertion (:attr:`ResidencyPlan.granularity_tax_bytes` + against :attr:`ResidencyPlan.leaf_granular_tax_bytes`). +* **The schedule** (:class:`ResidencySchedule`) — when the assigned {VRAM, RAM} + pair admits holding the whole ACTIVE component for a call, the cells rotate at + CALL boundaries (:meth:`StreamedResidency.hold_component`) instead of + per-forward. That is ``model_offload``'s per-CALL amortization — measured + 1.57× against this rung's per-STEP 1.91× at the same peak VRAM (pgw#1497) — + on this rung's mechanism and at this rung's finer granularity. The choice is + DETERMINISTIC ARITHMETIC over the budget and the cell layout, decided in the + plan, never a runtime heuristic reacting to pressure (Paul's boundary ruling). + +Admission-first is untouched: the budget is still handed down, the fill is still +the greedy largest-first walk run to a fixed point, and under the default +leaf-granular policy every number this module produces is bit-for-bit what it +produced before cells existed. """ from __future__ import annotations @@ -58,6 +86,7 @@ import logging import threading from dataclasses import dataclass +from enum import Enum from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple from . import staging @@ -78,6 +107,10 @@ #: weight's dtype, so every sub-buffer starts on a 512-byte boundary. _ALIGN = 512 +#: varena's VMM chunk granularity — the page size the driver maps in. The +#: number behind pgw#1507's measured 13.2 % tax, and the reason cells exist. +VMM_GRANULARITY_BYTES = 2 << 20 + #: Set on a module whose leaves this rung has hooked. ENGAGED_ATTR = "_cozy_stream_residency" @@ -86,8 +119,15 @@ _ADAPTER_MARKERS = ("lora_a", "lora_b", "lora_embedding", "lora_magnitude") +def align_to(n: int, alignment: int) -> int: + """``n`` rounded up to a multiple of ``alignment``. One rounding rule for + both alignments a mechanism has: the cast-buffer carve and the page.""" + a = max(1, int(alignment)) + return (int(n) + a - 1) // a * a + + def aligned(n: int) -> int: - return (int(n) + _ALIGN - 1) // _ALIGN * _ALIGN + return align_to(n, _ALIGN) # --------------------------------------------------------------------------- @@ -141,10 +181,214 @@ class LeafCost: name: str resident_bytes: int cast_bytes: int = 0 + #: Which COMPONENT this leaf belongs to — the first path segment of its + #: qualified name (``unet.down_blocks.0.attn.to_q`` -> ``unet``), which is + #: exactly the granularity ``model_offload`` swaps at and therefore the unit + #: the call-boundary schedule holds hot. A dotless name has no component and + #: joins the unnamed one, so a flat tree is one component rather than N. + component: str = "" def __post_init__(self) -> None: if not self.cast_bytes: object.__setattr__(self, "cast_bytes", int(self.resident_bytes)) + if not self.component and "." in self.name: + object.__setattr__(self, "component", self.name.split(".", 1)[0]) + + +@dataclass(frozen=True) +class CellPolicy: + """A MECHANISM's paging geometry — the only thing the planner needs to know + about how a rung actually holds bytes (pgw#1515). + + Three numbers, each a property of the mechanism and not of the model: + + * ``cell_bytes`` — the target size of one back/unback/stream unit. ``0`` + means LEAF-GRANULAR: every leaf is its own cell, which is what the rung + did before cells existed and what :data:`LEAF_CELLS` states. + * ``granularity_bytes`` — the page size a REGION is rounded up to. The + software cast ring has no region concept and states ``1``; varena maps in + 2 MiB VMM chunks and states :data:`VMM_GRANULARITY_BYTES`, which is the + whole source of pgw#1507's 13.2 % tax. + * ``intra_align_bytes`` — the alignment BETWEEN members inside one cell. A + uint8 slice must be re-viewable as its weight's dtype, so a mechanism that + carves views out of a shared region states :data:`_ALIGN`; one that does + not states ``1``. + + Stating a geometry is how a mechanism gets priced honestly: with + ``granularity_bytes=1`` the tax arithmetic collapses to zero and every + number the planner reports is the pre-cell number, exactly. + """ + + cell_bytes: int = 0 + granularity_bytes: int = 1 + intra_align_bytes: int = 1 + + @classmethod + def of(cls, value: "int | CellPolicy | None") -> "CellPolicy": + """Accept a bare target-cell-bytes int, a policy, or nothing.""" + if value is None: + return LEAF_CELLS + if isinstance(value, CellPolicy): + return value + return cls(cell_bytes=max(0, int(value))) + + @classmethod + def vmm(cls, cell_bytes: int) -> "CellPolicy": + """The varena geometry: 2 MiB pages, 512-byte views inside a cell.""" + return cls( + cell_bytes=max(0, int(cell_bytes)), + granularity_bytes=VMM_GRANULARITY_BYTES, + intra_align_bytes=_ALIGN, + ) + + def span_of(self, payload_bytes: int) -> int: + """What a region holding ``payload_bytes`` actually MAPS.""" + return align_to(payload_bytes, self.granularity_bytes) + + def member_bytes(self, leaf_bytes: int) -> int: + """What one member costs inside a cell, aligned for its neighbours.""" + return align_to(leaf_bytes, self.intra_align_bytes) + + +#: The default: one cell per leaf, no page granularity, no intra-cell padding. +#: Under this policy the planner's arithmetic is identical to pgw#1497's. +LEAF_CELLS = CellPolicy() + + +@dataclass(frozen=True) +class ResidencyCell: + """One back/unback/stream unit: a packed group of leaves that share a fate. + + Members of a cell are always in the same place. That is the point — a + mechanism that can map, stream or release half a cell is back to per-leaf + granularity and back to paying its page remainder per leaf. + """ + + index: int + component: str + members: Tuple[str, ...] + #: The weight bytes the members really hold (no alignment of any kind). + weight_bytes: int + #: What the mechanism MAPS for this cell: the members packed at the + #: intra-cell alignment, rounded up to one page granularity. + span_bytes: int + #: The in-flight cast cost when this cell streams — the whole cell, since + #: the cell is the stream unit. + cast_bytes: int + #: True when every member is always-resident (below the streaming floor, or + #: excluded by the caller). Forced and streamable leaves never share a cell: + #: a cell's fate has to be one thing. + forced: bool + #: The same members priced ONE REGION PER LEAF — the baseline this packing + #: exists to beat, carried so the elimination is reported and not asserted. + leaf_span_bytes: int + + @property + def granularity_tax_bytes(self) -> int: + return max(0, self.span_bytes - self.weight_bytes) + + @property + def leaf_granular_tax_bytes(self) -> int: + return max(0, self.leaf_span_bytes - self.weight_bytes) + + +def pack_cells( + costs: Sequence[LeafCost], + *, + policy: "int | CellPolicy | None" = None, + min_stream_bytes: int = DEFAULT_MIN_STREAM_BYTES, + exclude: Iterable[str] = (), +) -> Tuple[ResidencyCell, ...]: + """Group ``costs`` into cells. Deterministic, order-independent, pure. + + **Small leaves are packed FIRST**, which is not a tie-break but the whole + mechanism: a leaf's page remainder is bounded by the granularity, so a 4 KiB + norm weight in its own 2 MiB region wastes 99.8 % of it while a 300 MiB + attention block wastes at most 0.7 %. Filling ascending puts the small ones + together until a cell reaches its target and leaves the large ones as + singletons, where they were already efficient. pgw#1507 measured the cost of + not doing this: 248 MiB over 239 regions on a 1.830 GiB model. + + Two things a cell may never mix, because a cell has ONE fate and ONE home: + + * forced and streamable leaves — a forced member would pin the whole cell; + * two components — the call-boundary schedule holds a COMPONENT hot, so a + cell straddling two of them could not be swapped at a call boundary. + + Order-independence is a hard requirement inherited from + :func:`plan_residency`: groups are emitted in sorted key order and filled in + sorted member order, so the cells are a function of the SET of costs, never + of the sequence they arrived in. + """ + geometry = CellPolicy.of(policy) + skip = {str(n) for n in exclude} + floor = int(min_stream_bytes) + + groups: Dict[Tuple[str, bool], List[LeafCost]] = {} + for cost in costs: + forced = cost.name in skip or cost.cast_bytes < floor + groups.setdefault((cost.component, forced), []).append(cost) + + cells: List[ResidencyCell] = [] + for key in sorted(groups): + component, forced = key + members = sorted( + groups[key], key=lambda c: (c.cast_bytes, c.resident_bytes, c.name) + ) + batch: List[LeafCost] = [] + payload = 0 + for cost in members: + share = geometry.member_bytes(cost.resident_bytes) + if batch and payload + share > max(0, geometry.cell_bytes): + cells.append(_cell(len(cells), component, forced, batch, payload, geometry)) + batch, payload = [], 0 + batch.append(cost) + payload += share + if batch: + cells.append(_cell(len(cells), component, forced, batch, payload, geometry)) + return tuple(cells) + + +def _cell( + index: int, + component: str, + forced: bool, + members: Sequence[LeafCost], + payload_bytes: int, + geometry: CellPolicy, +) -> ResidencyCell: + return ResidencyCell( + index=index, + component=component, + members=tuple(c.name for c in members), + weight_bytes=sum(c.resident_bytes for c in members), + span_bytes=geometry.span_of(payload_bytes), + cast_bytes=sum(c.cast_bytes for c in members), + forced=forced, + leaf_span_bytes=sum( + geometry.span_of(geometry.member_bytes(c.resident_bytes)) for c in members + ), + ) + + +class ResidencySchedule(str, Enum): + """WHEN cells rotate — decided in the plan, never under fire. + + ``CALL_BOUNDARY`` is the ``model_offload``-equivalent mode on this rung's + mechanism: the assigned {VRAM, RAM} pair admits holding the whole active + component's cells for the duration of a call, so the swap happens once per + call and its cost amortises over every step (pgw#1497 measured that split + directly — ``model_offload`` 1.57× per-CALL against this rung's 1.91× + per-STEP at the SAME 1.9 GB peak). + + ``PER_STEP`` is the fallback the budget forces: no component fits whole, so + cells rotate inside the forward and the tax is paid every step. It is also + the mode that reaches below ``model_offload``'s floor, which is the band + this rung exists for. + """ + + CALL_BOUNDARY = "call_boundary" + PER_STEP = "per_step" @dataclass(frozen=True) @@ -169,6 +413,44 @@ class ResidencyPlan: #: enforced — see :class:`MemoryBudget`. ram_budget_bytes: int = 0 + # -- cells (pgw#1515) --------------------------------------------------- + #: The paging geometry this plan was drawn for. + policy: CellPolicy = LEAF_CELLS + #: Every cell, forced and streamable, in packing order. + cells: Tuple[ResidencyCell, ...] = () + #: The cells the card holds — indices into :attr:`cells`, in fill order. + resident_cells: Tuple[int, ...] = () + #: The cells that stream, in the same walk order. + streamed_cells: Tuple[int, ...] = () + #: What the mechanism actually MAPS for the resident set: the cell spans, + #: page-remainder included. Equals :attr:`resident_bytes` exactly when the + #: geometry has no page granularity, which is why the default policy leaves + #: every pre-cell number unchanged. + resident_span_bytes: int = 0 + #: The SAME resident leaves priced one region per leaf — pgw#1507's measured + #: baseline, carried so that :attr:`tax_eliminated_bytes` is a subtraction + #: of two numbers this plan computed rather than a claim about a past run. + leaf_granular_span_bytes: int = 0 + + # -- the schedule (pgw#1515) -------------------------------------------- + schedule: ResidencySchedule = ResidencySchedule.PER_STEP + #: (component, mapped span) for every component, largest first. The + #: arithmetic the schedule was decided on, published so the decision can be + #: audited without re-deriving it. + component_spans: Tuple[Tuple[str, int], ...] = () + #: The largest component's mapped span — what CALL_BOUNDARY must hold. + hot_component_bytes: int = 0 + #: What rests off-card while one component is hot. This is the RAM half's + #: real bill under CALL_BOUNDARY, and the reason a stated RAM budget can + #: force PER_STEP on a card that would otherwise admit the swap. + cold_component_bytes: int = 0 + #: The components that may be held whole for a call, largest first. Empty + #: under PER_STEP — an empty tuple IS the refusal. + hold_while_hot: Tuple[str, ...] = () + #: Which component this plan is currently holding hot, "" while the plan is + #: the budget-greedy split rather than a call-boundary arrangement. + hot_component: str = "" + @property def host_bytes(self) -> int: """What this plan costs in HOST RAM: the pinned streaming tail.""" @@ -183,8 +465,44 @@ def host_fits(self) -> bool: @property def device_bytes(self) -> int: - """What the card actually holds under this plan, at peak.""" - return self.resident_bytes + self.window_bytes + """What the card actually holds under this plan, at peak. + + The SPAN, not the weight bytes: a region is charged at what it maps, so + ``fits`` stays exact under a mechanism that pages (pgw#1507 charged its + regions the same way, which is what made its ``resident_bytes`` equal + the driver's mapped bytes to the byte).""" + return self.resident_span_bytes + self.window_bytes + + # -- the granularity tax, reported (pgw#1515) --------------------------- + + @property + def granularity_tax_bytes(self) -> int: + """What this plan's page remainders cost. Zero is the goal, and the + residual is published rather than assumed away.""" + return max(0, self.resident_span_bytes - self.resident_bytes) + + @property + def leaf_granular_tax_bytes(self) -> int: + """What the same resident leaves would have cost one region each — + pgw#1507 measured this at 248 MiB / 13.2 % on sd1.5.""" + return max(0, self.leaf_granular_span_bytes - self.resident_bytes) + + @property + def tax_eliminated_bytes(self) -> int: + """Weight bytes the packing bought back at the same lease.""" + return max(0, self.leaf_granular_span_bytes - self.resident_span_bytes) + + @property + def granularity_tax_ratio(self) -> float: + return self.granularity_tax_bytes / self.resident_bytes if self.resident_bytes else 0.0 + + @property + def leaf_granular_tax_ratio(self) -> float: + return ( + self.leaf_granular_tax_bytes / self.resident_bytes + if self.resident_bytes + else 0.0 + ) @property def fits(self) -> bool: @@ -205,6 +523,7 @@ def plan_residency( streams: int = DEFAULT_STREAMS, min_stream_bytes: int = DEFAULT_MIN_STREAM_BYTES, exclude: Iterable[str] = (), + cells: "int | CellPolicy | None" = None, ) -> ResidencyPlan: """Split ``costs`` into a resident set and a streaming tail under a budget. @@ -238,52 +557,126 @@ def plan_residency( always produce the same split. That is a hard requirement, not a nicety — a mint's traced graph specializations and a residency reservation both depend on the answer. + + **The unit of the walk is a CELL** (pgw#1515), not a leaf: cells are packed + first (:func:`pack_cells`), the fill admits whole cells charged at their + MAPPED SPAN, and the in-flight window reserves ``streams`` × the largest + streamed CELL, because a cell is what moves. Under the default + leaf-granular geometry every cell is one leaf with no page remainder, so + this is the identical arithmetic to the pre-cell planner — verified as a + standing test, not as an argument. + + The plan also decides its own SCHEDULE from the same two inputs, so that + "hold the hot component and swap at call boundaries" is a property of the + plan and never a runtime reaction to pressure. """ streams = max(1, int(streams)) pair = MemoryBudget.of(budget_bytes) budget = max(0, int(pair.vram_bytes)) - skip = {str(n) for n in exclude} - - forced: List[LeafCost] = [] - candidates: List[LeafCost] = [] - for cost in costs: - if cost.name in skip or cost.cast_bytes < int(min_stream_bytes): - forced.append(cost) - else: - candidates.append(cost) + geometry = CellPolicy.of(cells) + packed = pack_cells( + costs, policy=geometry, min_stream_bytes=min_stream_bytes, exclude=exclude + ) - order = sorted(candidates, key=lambda c: (-c.cast_bytes, -c.resident_bytes, c.name)) - floor = sum(c.resident_bytes for c in forced) + forced_cells = [c for c in packed if c.forced] + order = sorted( + (c for c in packed if not c.forced), + key=lambda c: (-c.cast_bytes, -c.weight_bytes, c.index), + ) + floor = sum(c.span_bytes for c in forced_cells) window = 0 - mem = floor - resident: List[LeafCost] = [] - streamed: List[LeafCost] = [] + span = floor + resident: List[ResidencyCell] = [] + streamed: List[ResidencyCell] = [] for _ in range(len(order) + 1): - mem = floor + span = floor resident = [] streamed = [] - for cost in order: - if mem + cost.resident_bytes + window <= budget: - resident.append(cost) - mem += cost.resident_bytes + for cell in order: + if span + cell.span_bytes + window <= budget: + resident.append(cell) + span += cell.span_bytes else: - streamed.append(cost) + streamed.append(cell) needed = streams * max((c.cast_bytes for c in streamed), default=0) if needed <= window: break window = needed + held = forced_cells + resident + component_span: Dict[str, int] = {} + for cell in packed: + component_span[cell.component] = component_span.get(cell.component, 0) + cell.span_bytes + ranked = tuple(sorted(component_span.items(), key=lambda kv: (-kv[1], kv[0]))) + hot = ranked[0][1] if ranked else 0 + cold = sum(v for _, v in ranked) - hot + # The swap moves a cell at a time and is pipelined exactly like a per-step + # cast, so the same ring reservation applies at a call boundary. + swap_window = streams * max((c.cast_bytes for c in packed), default=0) + ram = max(0, int(pair.ram_bytes)) + call_boundary = bool(ranked) and hot + swap_window <= budget and (not ram or cold <= ram) + return ResidencyPlan( budget_bytes=budget, streams=streams, - forced=tuple(c.name for c in forced), - resident=tuple(c.name for c in resident), - streamed=tuple(c.name for c in streamed), - resident_bytes=mem, - streamed_bytes=sum(c.resident_bytes for c in streamed), + forced=tuple(n for c in forced_cells for n in c.members), + resident=tuple(n for c in resident for n in c.members), + streamed=tuple(n for c in streamed for n in c.members), + resident_bytes=sum(c.weight_bytes for c in held), + streamed_bytes=sum(c.weight_bytes for c in streamed), window_bytes=window, - ram_budget_bytes=max(0, int(pair.ram_bytes)), + ram_budget_bytes=ram, + policy=geometry, + cells=packed, + resident_cells=tuple(c.index for c in held), + streamed_cells=tuple(c.index for c in streamed), + resident_span_bytes=span, + leaf_granular_span_bytes=sum(c.leaf_span_bytes for c in held), + schedule=( + ResidencySchedule.CALL_BOUNDARY if call_boundary else ResidencySchedule.PER_STEP + ), + component_spans=ranked, + hot_component_bytes=hot, + cold_component_bytes=cold, + hold_while_hot=tuple(name for name, _ in ranked) if call_boundary else (), + ) + + +def _hot_plan(plan: ResidencyPlan, component: str) -> ResidencyPlan: + """``plan`` rearranged so ``component`` is held whole and the rest parks. + + Pure, and derived from the cells the plan already carries — the schedule + decides WHETHER this arrangement is affordable, and this decides what it + looks like. The forced core (adapters, sub-floor leaves, caller exclusions) + stays resident in every arrangement: an exclusion is a statement about + residency, and a call boundary does not repeal it. + """ + held = [c for c in plan.cells if c.forced or c.component == component] + parked = [c for c in plan.cells if not c.forced and c.component != component] + window = plan.streams * max((c.cast_bytes for c in parked), default=0) + return ResidencyPlan( + budget_bytes=plan.budget_bytes, + streams=plan.streams, + forced=tuple(n for c in held if c.forced for n in c.members), + resident=tuple(n for c in held if not c.forced for n in c.members), + streamed=tuple(n for c in parked for n in c.members), + resident_bytes=sum(c.weight_bytes for c in held), + streamed_bytes=sum(c.weight_bytes for c in parked), + window_bytes=window, + ram_budget_bytes=plan.ram_budget_bytes, + policy=plan.policy, + cells=plan.cells, + resident_cells=tuple(c.index for c in held), + streamed_cells=tuple(c.index for c in parked), + resident_span_bytes=sum(c.span_bytes for c in held), + leaf_granular_span_bytes=sum(c.leaf_span_bytes for c in held), + schedule=plan.schedule, + component_spans=plan.component_spans, + hot_component_bytes=plan.hot_component_bytes, + cold_component_bytes=plan.cold_component_bytes, + hold_while_hot=plan.hold_while_hot, + hot_component=component, ) @@ -550,7 +943,15 @@ def discover_leaves( # reservation is the number the ring really allocates rather than # one that ignores the 512-byte sub-buffer boundaries. costs.append( - LeafCost(qualified, total, sum(aligned(tensor_bytes(t)) for _, _, t in own)) + LeafCost( + qualified, + total, + sum(aligned(tensor_bytes(t)) for _, _, t in own), + # STATED, not derived from the name: a root that is itself a + # leaf has a dotless qualified name and still belongs to its + # own component for the call-boundary schedule (pgw#1515). + component=root_name, + ) ) return leaves, costs, adapters @@ -585,6 +986,7 @@ def __init__( streams: int = DEFAULT_STREAMS, min_stream_bytes: int = DEFAULT_MIN_STREAM_BYTES, exclude: Iterable[str] = (), + cells: "int | CellPolicy | None" = None, ) -> None: import torch @@ -599,6 +1001,7 @@ def __init__( self.budget = MemoryBudget.of(budget_bytes) self.streams = max(1, int(streams)) self.min_stream_bytes = int(min_stream_bytes) + self.cells = CellPolicy.of(cells) self._exclude = {str(n) for n in exclude} self._roots = list(roots) self._leaves: Dict[str, Any] = {} @@ -653,6 +1056,7 @@ def engage(self) -> ResidencyPlan: streams=self.streams, min_stream_bytes=self.min_stream_bytes, exclude=self._exclude, + cells=self.cells, ), allow_promote=True, ) @@ -680,6 +1084,7 @@ def rebudget(self, budget_bytes: "int | MemoryBudget") -> ResidencyPlan: streams=self.streams, min_stream_bytes=self.min_stream_bytes, exclude=self._exclude, + cells=self.cells, ), allow_promote=allow_promote, ) @@ -703,6 +1108,49 @@ def partial_load(self, extra_bytes: int) -> int: after = self.plan.resident_bytes if self.plan is not None else 0 return max(0, after - before) + # -- the call-boundary schedule (pgw#1515) ------------------------------ + + def hold_component(self, component: str) -> ResidencyPlan: + """Hold ``component``'s cells whole and park every other component. + + THE call-boundary swap, and the ``model_offload``-equivalent mode on + this rung's mechanism: the hot component runs fully resident for the + whole call, so its move cost is paid once per CALL and amortises over + every step, instead of the per-STEP cast this rung otherwise pays. The + difference is measured, not argued — pgw#1497 clocked ``model_offload`` + at 1.57× against this rung's 1.91× at the SAME 1.9 GB peak, purely + because of where the tax lands. + + REFUSES when the plan's schedule is not + :attr:`ResidencySchedule.CALL_BOUNDARY`. A budget that cannot hold the + biggest component whole would answer this call by silently overshooting + its lease, and the whole point of deciding the schedule in the plan is + that the answer is known before anything moves. + """ + plan = self.plan if self.plan is not None else self.engage() + wanted = str(component) + if plan.schedule is not ResidencySchedule.CALL_BOUNDARY: + raise ValueError( + "hold_component(%r) refused: the assigned budget admits no whole " + "component (hot %d B + window %d B > VRAM %d B, cold %d B vs RAM %d B) " + "— this plan is scheduled %s" + % ( + wanted, + plan.hot_component_bytes, + plan.window_bytes, + plan.budget_bytes, + plan.cold_component_bytes, + plan.ram_budget_bytes, + plan.schedule.value, + ) + ) + if wanted not in plan.hold_while_hot: + raise ValueError( + "hold_component(%r) refused: no such component — this tree has %r" + % (wanted, plan.hold_while_hot) + ) + return self._apply(_hot_plan(plan, wanted), allow_promote=True) + def demote_to_host(self) -> int: """Every weight to pinned host RAM. The residency HOST tier.""" return self.partial_unload(self.total_bytes) @@ -757,6 +1205,11 @@ def _recorded( streamed = tuple(old.streamed) + tuple( n for n in old.all_resident if n in dropped ) + # The cells that survived the trim, so the span and tax numbers keep + # describing what is really mapped rather than what the planner drew. + kept = set(forced) | set(resident) + held = [c for c in planned.cells if kept.issuperset(c.members)] + parked = [c for c in planned.cells if not kept.issuperset(c.members)] return ResidencyPlan( budget_bytes=planned.budget_bytes, streams=planned.streams, @@ -771,6 +1224,17 @@ def _recorded( ), window_bytes=planned.window_bytes, ram_budget_bytes=planned.ram_budget_bytes, + policy=planned.policy, + cells=planned.cells, + resident_cells=tuple(c.index for c in held), + streamed_cells=tuple(c.index for c in parked), + resident_span_bytes=sum(c.span_bytes for c in held), + leaf_granular_span_bytes=sum(c.leaf_span_bytes for c in held), + schedule=planned.schedule, + component_spans=planned.component_spans, + hot_component_bytes=planned.hot_component_bytes, + cold_component_bytes=planned.cold_component_bytes, + hold_while_hot=planned.hold_while_hot, ) def _place_residue(self) -> None: @@ -938,11 +1402,17 @@ def stream_residency_active(obj: Any) -> bool: "DEFAULT_MIN_STREAM_BYTES", "DEFAULT_STREAMS", "ENGAGED_ATTR", + "LEAF_CELLS", + "VMM_GRANULARITY_BYTES", + "CellPolicy", "LeafCost", "MemoryBudget", "PlanTransition", + "ResidencyCell", "ResidencyPlan", + "ResidencySchedule", "StreamedResidency", + "align_to", "aligned", "bind_tensor", "discover_leaves", @@ -950,6 +1420,7 @@ def stream_residency_active(obj: Any) -> bool: "is_streamable_leaf", "module_roots", "own_tensors", + "pack_cells", "tensor_bytes", "plan_residency", "plan_transition", diff --git a/tests/test_residency_cells.py b/tests/test_residency_cells.py new file mode 100644 index 000000000..dcdaa81ba --- /dev/null +++ b/tests/test_residency_cells.py @@ -0,0 +1,596 @@ +"""Coarse residency cells and the hold-while-hot schedule. + +# pgw#1515: the packing-vs-scheduling split, phase 3's residency half. + +Two claims are under test, and both are ARITHMETIC the plan publishes rather +than assertions about a past run: + +1. **Packing eliminates the per-leaf granularity tax.** pgw#1507 measured it on + the card: 1.830 GiB of sd1.5 weights held in 2.072 GiB of span across 239 + per-leaf regions — 248 MiB, **13.2 %**, forced by 2 MiB VMM granularity plus + independent per-leaf back/unback. Every test below computes BOTH prices from + the same leaf census, so the elimination is a subtraction and the baseline + can be checked against the measured number. +2. **The schedule is decided in the plan.** A budget that admits holding the + whole active component for a call gets ``CALL_BOUNDARY`` — the per-CALL + amortization ``model_offload`` gets, on this rung's mechanism — and one that + does not gets ``PER_STEP``. Deterministic from budget + cell layout, never a + runtime reaction to pressure. + +Real ``nn.Module`` trees throughout for anything touching the mechanism; the +planner tests feed it :class:`LeafCost`, which is its real input type, not a +mock of one. CPU only: the cell layer is arithmetic plus the same host/device +mover the CPU arm already exercises. +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +# A real static import, not `nn = torch.nn` — the rebinding form leaves mypy no +# base type for `class X(nn.Module)`. Same idiom as `test_stream_residency.py`. +import torch.nn as nn # noqa: E402 + +from gen_worker.models.stream_residency import ( # noqa: E402 + LEAF_CELLS, + VMM_GRANULARITY_BYTES, + CellPolicy, + LeafCost, + ResidencySchedule, + StreamedResidency, + module_roots, + pack_cells, + plan_residency, +) + +MIB = 1 << 20 + + +# --------------------------------------------------------------------------- +# Fixtures: a real two-component tree, and a leaf census shaped like the one +# pgw#1507 measured. +# --------------------------------------------------------------------------- + + +class Block(nn.Module): + def __init__(self, width: int) -> None: + super().__init__() + self.fc1 = nn.Linear(width, width * 2) + self.fc2 = nn.Linear(width * 2, width) + self.norm = nn.LayerNorm(width) + + def forward(self, x): # type: ignore[no-untyped-def] + return self.norm(x + self.fc2(torch.relu(self.fc1(x)))) + + +class Tower(nn.Module): + def __init__(self, width: int, depth: int) -> None: + super().__init__() + self.blocks = nn.ModuleList([Block(width) for _ in range(depth)]) + + def forward(self, x): # type: ignore[no-untyped-def] + for block in self.blocks: + x = block(x) + return x + + +class TwoComponent: + """A pipeline-shaped holder: two independent module trees under one object. + + ``module_roots`` reads exactly this shape off the serve path (an author + model object whose attributes are modules), so the components the schedule + reasons about are the ones the real walk produces. + """ + + def __init__(self, width: int = 192, depth: int = 4) -> None: + torch.manual_seed(1515) + self.unet = Tower(width, depth).eval() + self.text_encoder = Tower(width // 2, max(1, depth // 2)).eval() + + +def _census(n: int = 239, total_bytes: int = int(1.830 * (1 << 30))) -> list[LeafCost]: + """A leaf census with pgw#1507's shape: 239 leaves, 1.830 GiB, remainders + spread across the 2 MiB granularity so the per-leaf tax lands in the + measured band. Deterministic — a fixed LCG, no RNG dependency.""" + sizes: list[int] = [] + state = 1515 + for _ in range(n): + state = (state * 1103515245 + 12345) % (1 << 31) + # fp16 tensors: even byte counts, sizes spread over two orders of + # magnitude the way an attention tower's leaves are. + sizes.append(2 * (1 + state % (4 * MIB))) + scale = total_bytes / sum(sizes) + return [ + LeafCost(f"unet.leaf{i:03d}", max(2, 2 * int(size * scale / 2))) + for i, size in enumerate(sizes) + ] + + +def _vmm(cell_bytes: int) -> CellPolicy: + return CellPolicy.vmm(cell_bytes) + + +# --------------------------------------------------------------------------- +# 1. Packing +# --------------------------------------------------------------------------- + + +def test_packing_is_deterministic_under_input_order() -> None: + """The cells must be a function of the SET of leaves, never the sequence + they arrived in — a mint's graph specializations and a residency + reservation both key on the answer.""" + costs = _census(64) + first = pack_cells(costs, policy=_vmm(64 * MIB), min_stream_bytes=1) + shuffled = pack_cells( + list(reversed(costs)), policy=_vmm(64 * MIB), min_stream_bytes=1 + ) + assert first == shuffled + assert len(first) < len(costs), "64 MiB cells must actually pack this census" + + +def test_packing_eliminates_the_per_leaf_granularity_tax() -> None: + """THE claim. Both prices come out of the same census, so the baseline is + checkable against pgw#1507's measured 13.2 % and the residual is a number, + not a promise. + + The residual is NOT zero and is not claimed to be: a cell still pays one + page remainder, so the tax falls as ``cells x granularity / 2`` — 244 MiB + over 239 leaf regions becomes 32 MiB over 33 cells at 64 MiB and 6 MiB over + 8 cells at 256 MiB. What is eliminated is the PER-LEAF term.""" + costs = _census() + weights = sum(c.resident_bytes for c in costs) + + leafwise = pack_cells(costs, policy=_vmm(0), min_stream_bytes=1) + leaf_span = sum(c.span_bytes for c in leafwise) + leaf_tax = (leaf_span - weights) / weights + assert len(leafwise) == len(costs), "leaf-granular means one region per leaf" + # The measured baseline this packing exists to beat: 13.2 % on 239 regions. + assert 0.10 <= leaf_tax <= 0.16, f"census is off the measured shape: {leaf_tax:.3%}" + + for target, ceiling in ((64 * MIB, 0.02), (256 * MIB, 0.005)): + packed = pack_cells(costs, policy=_vmm(target), min_stream_bytes=1) + packed_span = sum(c.span_bytes for c in packed) + packed_tax = (packed_span - weights) / weights + assert packed_tax < ceiling, f"{target // MIB} MiB cells: {packed_tax:.3%}" + assert packed_span < leaf_span + # BOUNDED, not merely small: at most one page remainder per cell. + assert packed_span - weights <= len(packed) * VMM_GRANULARITY_BYTES + + +def test_the_residual_tax_is_bounded_by_the_geometry_not_by_the_leaves() -> None: + """The exact shape of the elimination, and the reason it is a DESIGN and not + a heuristic: the leaf-granular tax is unbounded as leaves shrink (a 64 KiB + leaf wastes 97 % of its 2 MiB page, and sd1.5's census is full of them), + while the packed residual can never exceed ``granularity / cell_bytes`` — + 3.1 % at 64 MiB cells, 0.8 % at 256 MiB — whatever the leaves look like.""" + costs = [LeafCost(f"unet.s{i:04d}", 64 * 1024 + 2 * i) for i in range(4096)] + weights = sum(c.resident_bytes for c in costs) + + leafwise = sum( + c.span_bytes for c in pack_cells(costs, policy=_vmm(0), min_stream_bytes=1) + ) + leaf_tax = (leafwise - weights) / weights + assert leaf_tax > 20, "a 64 KiB leaf wastes a whole 2 MiB page" + + for target in (64 * MIB, 256 * MIB): + packed = pack_cells(costs, policy=_vmm(target), min_stream_bytes=1) + span = sum(c.span_bytes for c in packed) + residual = (span - weights) / weights + assert residual <= VMM_GRANULARITY_BYTES / target + assert residual < leaf_tax / 500 + + +def test_bigger_cells_monotonically_shrink_the_residual_tax() -> None: + """The sizing knob has to be a real dial, not a switch — tomorrow's pricing + sweep is {leaf, 64 MiB, 256 MiB, component} and a non-monotone tax would + make that table unreadable.""" + costs = _census() + weights = sum(c.resident_bytes for c in costs) + taxes = [ + sum(c.span_bytes for c in pack_cells(costs, policy=_vmm(size), min_stream_bytes=1)) + - weights + for size in (0, 16 * MIB, 64 * MIB, 256 * MIB, 1 << 30) + ] + assert taxes == sorted(taxes, reverse=True) + assert taxes[0] > 100 * MIB and taxes[-1] < 4 * MIB + + +def test_small_leaves_are_packed_first() -> None: + """Packing small-first is the mechanism, not a tie-break: a 4 KiB norm in + its own 2 MiB region wastes 99.8 % of it, a 300 MiB block wastes 0.7 %.""" + small = [LeafCost(f"unet.s{i}", 64 * 1024) for i in range(40)] + large = [LeafCost(f"unet.big{i}", 100 * MIB) for i in range(2)] + cells = pack_cells(small + large, policy=_vmm(8 * MIB), min_stream_bytes=1) + shared = [c for c in cells if len(c.members) > 1] + singles = [c for c in cells if len(c.members) == 1] + assert shared, "the small leaves must have been combined" + assert all(m.startswith("unet.s") for c in shared for m in c.members) + assert {c.members[0] for c in singles} == {"unet.big0", "unet.big1"} + + +def test_a_cell_never_mixes_forced_and_streamable_leaves() -> None: + """A cell has ONE fate. A forced member would pin the whole cell onto the + card, which is how a packing layer silently un-does a budget.""" + costs = [LeafCost(f"unet.l{i}", (i + 1) * MIB) for i in range(8)] + cells = pack_cells( + costs, policy=_vmm(64 * MIB), min_stream_bytes=1, exclude=("unet.l7",) + ) + for cell in cells: + assert cell.forced == ("unet.l7" in cell.members) or "unet.l7" not in cell.members + forced_cells = [c for c in cells if c.forced] + assert [c.members for c in forced_cells] == [("unet.l7",)] + + +def test_a_cell_never_straddles_two_components() -> None: + """The call-boundary schedule swaps a COMPONENT; a cell spanning two of + them could not be moved at a call boundary at all.""" + costs = [LeafCost(f"unet.l{i}", MIB) for i in range(4)] + [ + LeafCost(f"vae.l{i}", MIB) for i in range(4) + ] + cells = pack_cells(costs, policy=_vmm(64 * MIB), min_stream_bytes=1) + for cell in cells: + assert len({m.split(".", 1)[0] for m in cell.members}) == 1 + assert all(m.startswith(cell.component + ".") for m in cell.members) + + +# --------------------------------------------------------------------------- +# 2. The plan reports the tax, and the fill admits whole cells +# --------------------------------------------------------------------------- + + +def test_the_plan_reports_the_tax_it_did_not_pay() -> None: + costs = _census() + budget = int(1.2 * (1 << 30)) + leafwise = plan_residency( + costs, budget_bytes=budget, min_stream_bytes=1, cells=_vmm(0) + ) + packed = plan_residency( + costs, budget_bytes=budget, min_stream_bytes=1, cells=_vmm(64 * MIB) + ) + assert leafwise.leaf_granular_tax_ratio == pytest.approx( + leafwise.granularity_tax_ratio + ), "leaf-granular IS the baseline, so it can eliminate nothing" + assert leafwise.tax_eliminated_bytes == 0 + + # Measured by this plan, on this census: 8.0 % leaf-granular -> 1.6 % packed. + assert leafwise.granularity_tax_ratio > 0.05 + assert packed.granularity_tax_ratio < 0.02 + assert packed.leaf_granular_tax_ratio > 0.05 + assert packed.granularity_tax_ratio < packed.leaf_granular_tax_ratio / 5 + assert packed.tax_eliminated_bytes > 0 + + +def test_coarser_cells_buy_tax_and_pay_window_and_the_plan_prices_both() -> None: + """The counter-force, and it is NOT small — banked here because the GPU + sweep has to price it and a planner that only reported the win would hide + the reason a bigger cell can be worse. + + A cell is the stream unit, so under PER_STEP the in-flight reservation is + ``streams x the largest streamed CELL``: at a 1.2 GiB budget on this census + the window goes 20 MiB (leaf) -> 115 MiB (64 MiB cells) -> 508 MiB (256 MiB + cells), and it comes straight out of the resident set. So the tax falls + monotonically while the WEIGHT HELD peaks at a moderate cell size — 1118 / + 1112 / 1086 / 630 MiB across those four geometries. Under CALL_BOUNDARY + nothing streams mid-forward and the elimination is pure gain; under PER_STEP + it is a trade, and tomorrow's {leaf, 64 MiB, 256 MiB, component} x budget + table is the measurement of exactly this curve.""" + costs = _census() + budget = int(1.2 * (1 << 30)) + plans = [ + plan_residency(costs, budget_bytes=budget, min_stream_bytes=1, cells=_vmm(size)) + for size in (0, 16 * MIB, 64 * MIB, 256 * MIB) + ] + windows = [p.window_bytes for p in plans] + taxes = [p.granularity_tax_ratio for p in plans] + assert windows == sorted(windows), "a coarser cell reserves a bigger window" + assert taxes[-1] < taxes[0] / 10, "and buys the tax back" + assert plans[-1].resident_bytes < plans[0].resident_bytes, ( + "at 256 MiB the window costs more than the tax it saves" + ) + + +def test_the_fill_charges_a_cell_its_mapped_span_not_its_weight() -> None: + """``fits`` has to stay exact under a mechanism that pages — pgw#1507's + numbers matched the driver's mapped bytes precisely because regions were + priced at their aligned span.""" + costs = [LeafCost(f"unet.l{i}", 3 * MIB + 1024) for i in range(8)] + plan = plan_residency( + costs, budget_bytes=20 * MIB, min_stream_bytes=1, cells=_vmm(0) + ) + held = [c for c in plan.cells if c.index in plan.resident_cells] + assert plan.resident_span_bytes == sum(c.span_bytes for c in held) + assert plan.resident_span_bytes > plan.resident_bytes, "3 MiB + 1 KiB maps 4 MiB" + assert plan.device_bytes == plan.resident_span_bytes + plan.window_bytes + assert plan.fits and plan.device_bytes <= plan.budget_bytes + + +def test_a_cells_members_share_one_fate() -> None: + costs = _census(80) + plan = plan_residency( + costs, + budget_bytes=int(0.4 * (1 << 30)), + min_stream_bytes=1, + cells=_vmm(64 * MIB), + ) + assert plan.resident and plan.streamed, "this budget must genuinely split" + resident, streamed = set(plan.all_resident), set(plan.streamed) + for cell in plan.cells: + members = set(cell.members) + assert members <= resident or members <= streamed + + +def test_the_window_reserves_the_largest_streamed_CELL() -> None: + """A cell is what moves, so the in-flight reservation grows with the cell + size. Coarser cells are not free and the plan says so.""" + costs = [LeafCost(f"unet.l{i}", 4 * MIB) for i in range(16)] + plan = plan_residency( + costs, budget_bytes=40 * MIB, streams=2, min_stream_bytes=1, cells=_vmm(16 * MIB) + ) + assert plan.streamed_cells + largest = max(plan.cells[i].cast_bytes for i in plan.streamed_cells) + assert plan.window_bytes == 2 * largest + assert plan.device_bytes <= plan.budget_bytes + + +def test_the_default_policy_is_the_pre_cell_planner_exactly() -> None: + """The reduction that makes this a safe layer to add: state no geometry and + every number is the one pgw#1497 shipped.""" + costs = [LeafCost(f"unet.l{i}", (i + 1) * 1_000_000) for i in range(6)] + plan = plan_residency(costs, budget_bytes=14_000_000, min_stream_bytes=1) + assert plan.policy == LEAF_CELLS + assert len(plan.cells) == len(costs) + assert all(len(c.members) == 1 for c in plan.cells) + assert plan.resident_span_bytes == plan.resident_bytes + assert plan.granularity_tax_bytes == 0 and plan.leaf_granular_tax_bytes == 0 + assert plan.device_bytes == plan.resident_bytes + plan.window_bytes + + +def test_the_vram_ram_pair_shape_survives_the_cell_layer() -> None: + """pgw#1497's pair is the plan's shape and cells do not narrow it.""" + from gen_worker.models.stream_residency import MemoryBudget + + costs = _census(64) + tail = 200 * MIB + plan = plan_residency( + costs, + budget_bytes=MemoryBudget(vram_bytes=300 * MIB, ram_bytes=tail), + min_stream_bytes=1, + cells=_vmm(64 * MIB), + ) + assert plan.ram_budget_bytes == tail + assert plan.host_bytes == plan.streamed_bytes + assert plan.host_fits is (plan.host_bytes <= tail) + unstated = plan_residency( + costs, budget_bytes=300 * MIB, min_stream_bytes=1, cells=_vmm(64 * MIB) + ) + assert unstated.ram_budget_bytes == 0 and unstated.host_fits + + +# --------------------------------------------------------------------------- +# 3. The schedule +# --------------------------------------------------------------------------- + + +def _two_components(unet_mb: int = 400, text_mb: int = 100) -> list[LeafCost]: + return [LeafCost(f"unet.l{i}", 4 * MIB) for i in range(unet_mb // 4)] + [ + LeafCost(f"text_encoder.l{i}", 4 * MIB) for i in range(text_mb // 4) + ] + + +def test_a_budget_that_holds_the_hot_component_schedules_at_call_boundaries() -> None: + costs = _two_components() + plan = plan_residency( + costs, budget_bytes=600 * MIB, streams=2, min_stream_bytes=1, cells=_vmm(64 * MIB) + ) + assert plan.schedule is ResidencySchedule.CALL_BOUNDARY + assert plan.hold_while_hot == ("unet", "text_encoder") + assert plan.hot_component_bytes + plan.streams * max( + c.cast_bytes for c in plan.cells + ) <= plan.budget_bytes + + +def test_a_budget_below_the_hot_component_falls_to_per_step() -> None: + costs = _two_components() + plan = plan_residency( + costs, budget_bytes=200 * MIB, streams=2, min_stream_bytes=1, cells=_vmm(64 * MIB) + ) + assert plan.schedule is ResidencySchedule.PER_STEP + assert plan.hold_while_hot == () + assert plan.hot_component_bytes > plan.budget_bytes + + +def test_the_schedule_flips_exactly_where_the_arithmetic_says() -> None: + """Deterministic from budget + layout: one byte on either side of the + stated inequality, and nothing else moves.""" + costs = _two_components() + probe = plan_residency( + costs, budget_bytes=1 << 40, streams=2, min_stream_bytes=1, cells=_vmm(64 * MIB) + ) + threshold = probe.hot_component_bytes + probe.streams * max( + c.cast_bytes for c in probe.cells + ) + at = plan_residency( + costs, budget_bytes=threshold, streams=2, min_stream_bytes=1, cells=_vmm(64 * MIB) + ) + below = plan_residency( + costs, budget_bytes=threshold - 1, streams=2, min_stream_bytes=1, + cells=_vmm(64 * MIB), + ) + assert at.schedule is ResidencySchedule.CALL_BOUNDARY + assert below.schedule is ResidencySchedule.PER_STEP + + +def test_a_stated_ram_half_can_force_per_step_on_a_card_that_would_admit_it() -> None: + """The pair is one decision, not two. VRAM admitting the hot component is + worthless if the RAM half cannot hold what parks while it is hot.""" + from gen_worker.models.stream_residency import MemoryBudget + + costs = _two_components() + roomy = plan_residency( + costs, + budget_bytes=MemoryBudget(600 * MIB, 1 << 30), + streams=2, + min_stream_bytes=1, + cells=_vmm(64 * MIB), + ) + assert roomy.schedule is ResidencySchedule.CALL_BOUNDARY + + starved = plan_residency( + costs, + budget_bytes=MemoryBudget(600 * MIB, 8 * MIB), + streams=2, + min_stream_bytes=1, + cells=_vmm(64 * MIB), + ) + assert starved.cold_component_bytes > 8 * MIB + assert starved.schedule is ResidencySchedule.PER_STEP + + +def test_the_schedule_is_deterministic_and_component_spans_are_published() -> None: + costs = _two_components() + first = plan_residency( + costs, budget_bytes=600 * MIB, min_stream_bytes=1, cells=_vmm(64 * MIB) + ) + again = plan_residency( + list(reversed(costs)), budget_bytes=600 * MIB, min_stream_bytes=1, + cells=_vmm(64 * MIB), + ) + assert first == again + assert [name for name, _ in first.component_spans] == ["unet", "text_encoder"] + assert first.hot_component_bytes == first.component_spans[0][1] + assert first.cold_component_bytes == sum(v for _, v in first.component_spans[1:]) + + +# --------------------------------------------------------------------------- +# 4. Holding a component hot, over a real tree +# --------------------------------------------------------------------------- + + +def test_holding_a_component_hot_parks_every_other_component( +) -> None: + model = TwoComponent() + roots = module_roots(model) + assert {name for name, _ in roots} == {"unet", "text_encoder"} + residency = StreamedResidency( + roots, device="cpu", budget_bytes=1 << 40, min_stream_bytes=1, + cells=CellPolicy.vmm(4 * MIB), + ) + plan = residency.engage() + assert plan.schedule is ResidencySchedule.CALL_BOUNDARY + + hot = residency.hold_component("unet") + assert hot.hot_component == "unet" + assert all(n.startswith("unet.") for n in hot.resident) + assert all(n.startswith("text_encoder.") for n in hot.streamed) + assert hot.streamed, "the cold component must actually park" + + cold = residency.hold_component("text_encoder") + assert all(n.startswith("text_encoder.") for n in cold.resident) + assert all(n.startswith("unet.") for n in cold.streamed) + + +def test_a_held_component_computes_the_same_answer_as_a_resident_one() -> None: + """The swap is a placement, never a numeric change — the property the whole + rung rests on, re-checked across a call-boundary rotation.""" + model = TwoComponent() + x = torch.randn(2, 192) + y = torch.randn(2, 96) + with torch.no_grad(): + want_unet = model.unet(x).clone() + want_text = model.text_encoder(y).clone() + + residency = StreamedResidency( + module_roots(model), device="cpu", budget_bytes=1 << 40, min_stream_bytes=1, + cells=CellPolicy.vmm(4 * MIB), + ) + residency.engage() + for _ in range(2): + residency.hold_component("unet") + with torch.no_grad(): + assert torch.equal(model.unet(x), want_unet) + assert torch.equal(model.text_encoder(y), want_text) + residency.hold_component("text_encoder") + with torch.no_grad(): + assert torch.equal(model.unet(x), want_unet) + assert torch.equal(model.text_encoder(y), want_text) + residency.release() + with torch.no_grad(): + assert torch.equal(model.unet(x), want_unet) + + +def test_hold_component_refuses_when_the_budget_cannot_admit_it() -> None: + """A refusal, never a silent overshoot: the plan already knows the answer + before a byte moves, which is the whole reason the schedule is planned.""" + model = TwoComponent() + residency = StreamedResidency( + module_roots(model), device="cpu", budget_bytes=64 * 1024, min_stream_bytes=1, + cells=CellPolicy.vmm(4 * MIB), + ) + plan = residency.engage() + assert plan.schedule is ResidencySchedule.PER_STEP + with pytest.raises(ValueError, match="admits no whole component"): + residency.hold_component("unet") + + +def test_hold_component_refuses_a_component_this_tree_does_not_have() -> None: + model = TwoComponent() + residency = StreamedResidency( + module_roots(model), device="cpu", budget_bytes=1 << 40, min_stream_bytes=1, + cells=CellPolicy.vmm(4 * MIB), + ) + residency.engage() + with pytest.raises(ValueError, match="no such component"): + residency.hold_component("vae") + + +def test_the_forced_core_stays_resident_through_a_call_boundary_swap() -> None: + """An exclusion is a statement about RESIDENCY (pgw#1497 defect 2), and a + call boundary does not repeal it — LoRA adapters live in this set.""" + model = TwoComponent() + roots = module_roots(model) + residency = StreamedResidency( + roots, + device="cpu", + budget_bytes=1 << 40, + min_stream_bytes=1, + exclude=("text_encoder.blocks.0.fc1",), + cells=CellPolicy.vmm(4 * MIB), + ) + residency.engage() + hot = residency.hold_component("unet") + assert "text_encoder.blocks.0.fc1" in hot.forced + assert "text_encoder.blocks.0.fc1" not in hot.streamed + param = model.text_encoder.blocks[0].fc1.weight # type: ignore[index,union-attr] + assert not param.is_pinned() + + +def test_a_real_tree_reports_its_own_tax_elimination() -> None: + """End to end on a real ``nn.Module`` census: the same tree, two geometries, + and the plan's own numbers show what the packing bought.""" + # 512-wide, 8 deep: ~33 MB of real parameters, enough that a 2 MiB page is + # a rounding detail rather than the whole model. + model = TwoComponent(width=512, depth=8) + roots = module_roots(model) + leafwise = StreamedResidency( + roots, device="cpu", budget_bytes=1 << 40, min_stream_bytes=1, + cells=CellPolicy.vmm(0), + ) + packed = StreamedResidency( + roots, device="cpu", budget_bytes=1 << 40, min_stream_bytes=1, + cells=CellPolicy.vmm(64 * MIB), + ) + a = leafwise.engage() + b = packed.engage() + assert a.resident_bytes == b.resident_bytes, "same weights, different geometry" + assert a.granularity_tax_ratio > 0.5, "small leaves in 2 MiB regions are brutal" + # The claim is the SHAPE: a per-CELL remainder where there was a per-LEAF + # one. On a 33 MB tree two cells' remainders are still ~10 % of it, which is + # the geometry being honest about a model smaller than a few pages. + held = [c for c in b.cells if c.index in b.resident_cells] + assert b.granularity_tax_bytes <= len(held) * VMM_GRANULARITY_BYTES + assert b.granularity_tax_ratio < a.granularity_tax_ratio / 5 + assert b.tax_eliminated_bytes > 0 + packed.release() + leafwise.release() From 8a53986d81788887fa72b254b194b4d49f0e48d1 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Wed, 19 Aug 2026 18:50:09 -0600 Subject: [PATCH 2/2] pgw#1515: the swap ring reserves only for cells that can MOVE, and the tests that were not discriminating MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three things the red-arm audit and the gate found: * The call-boundary reservation was streams x the largest cell of ANY kind, including the forced core — which packing makes large (500 sub-floor leaves become ONE cell). A cell that is resident in every arrangement never travels, so reserving the ring for it denied CALL_BOUNDARY to budgets that plainly admit it. Caught by a new test, red 5/5 under the old arithmetic. * test_small_leaves_are_packed_first passed under a REVERSED sort (audit M1, missed 0/5) — it asserted an outcome greedy fill produces either way. Replaced with the outcome small-first actually buys: a leaf at or above the cell target lands ALONE, so small leaves are never welded to a big one's residency. Red 5/5 under the reversal. * hold_component has no production caller yet — the serve path has no call-BOUNDARY hook and adding one is the half that must be measured. Baselined in scripts/unreached_surface_baseline.txt with an owner and an expiry (the pgw#1515 GPU sweep either produces the caller or the deletion), not left to go red on master. --- scripts/unreached_surface_baseline.txt | 22 +++++++ src/gen_worker/models/stream_residency.py | 9 ++- tests/test_residency_cells.py | 72 ++++++++++++++++++----- 3 files changed, 86 insertions(+), 17 deletions(-) diff --git a/scripts/unreached_surface_baseline.txt b/scripts/unreached_surface_baseline.txt index 7d4a6eef0..8a295dd62 100644 --- a/scripts/unreached_surface_baseline.txt +++ b/scripts/unreached_surface_baseline.txt @@ -527,3 +527,25 @@ gen_worker.serving.engine_runtime.process_vram_bytes() # caller and the ratchet will say so by going red on STALE. # EXPIRY: none. A guard's undo living beside its do is the correct shape. gen_worker.serving.weightless_program.uninstall() + +# Built ahead of its caller, on purpose, with the caller NAMED +# StreamedResidency.hold_component — pgw#1515's call-boundary swap. The +# PLAN already decides the schedule (ResidencySchedule.CALL_BOUNDARY vs +# PER_STEP) from the assigned {VRAM, RAM} pair and the cell layout, and this +# is the method that executes it: hold one component whole for a call, park +# the others. It is unreached because the serve path has no call-BOUNDARY +# hook yet — the rung engages once per placement, not once per generation — +# and adding one is the half that must be MEASURED, not reasoned about +# (pgw#1497 measured the per-CALL vs per-STEP split at 1.57x vs 1.91x on the +# card; whether our finer-grained version of it beats model_offload at its +# own 1.9 GB floor is the open question this whole issue exists to answer). +# Wiring it before that measurement would put a rung on the ladder at a +# price nobody has paid, which is precisely what pgw#1507 refused to do. +# +# OWNER: pgw#1515. WIRE IT — this is not a permanent exemption. The GPU +# pricing sweep ({leaf, 64 MiB, 256 MiB, component} x four budgets) +# either produces the caller or produces the deletion. +# EXPIRY: the pgw#1515 GPU sweep. If that sweep says the schedule does not +# beat model_offload at its floor, this method and its tests go, +# they are not ported (DESIGN-RULINGS 4.34). +gen_worker.models.stream_residency.StreamedResidency.hold_component() diff --git a/src/gen_worker/models/stream_residency.py b/src/gen_worker/models/stream_residency.py index 874fa859c..c756e5139 100644 --- a/src/gen_worker/models/stream_residency.py +++ b/src/gen_worker/models/stream_residency.py @@ -612,8 +612,13 @@ def plan_residency( hot = ranked[0][1] if ranked else 0 cold = sum(v for _, v in ranked) - hot # The swap moves a cell at a time and is pipelined exactly like a per-step - # cast, so the same ring reservation applies at a call boundary. - swap_window = streams * max((c.cast_bytes for c in packed), default=0) + # cast, so the same ring reservation applies at a call boundary — but over + # the cells that can actually MOVE. Forced cells are resident in every + # arrangement, and packing makes that distinction matter: a tree with 500 + # sub-floor leaves has ONE large forced cell, and reserving the ring for a + # cell that never travels would deny the call-boundary schedule to budgets + # that plainly admit it. + swap_window = streams * max((c.cast_bytes for c in packed if not c.forced), default=0) ram = max(0, int(pair.ram_bytes)) call_boundary = bool(ranked) and hot + swap_window <= budget and (not ram or cold <= ram) diff --git a/tests/test_residency_cells.py b/tests/test_residency_cells.py index dcdaa81ba..539a89414 100644 --- a/tests/test_residency_cells.py +++ b/tests/test_residency_cells.py @@ -197,17 +197,28 @@ def test_bigger_cells_monotonically_shrink_the_residual_tax() -> None: assert taxes[0] > 100 * MIB and taxes[-1] < 4 * MIB -def test_small_leaves_are_packed_first() -> None: - """Packing small-first is the mechanism, not a tie-break: a 4 KiB norm in - its own 2 MiB region wastes 99.8 % of it, a 300 MiB block wastes 0.7 %.""" - small = [LeafCost(f"unet.s{i}", 64 * 1024) for i in range(40)] - large = [LeafCost(f"unet.big{i}", 100 * MIB) for i in range(2)] - cells = pack_cells(small + large, policy=_vmm(8 * MIB), min_stream_bytes=1) - shared = [c for c in cells if len(c.members) > 1] - singles = [c for c in cells if len(c.members) == 1] - assert shared, "the small leaves must have been combined" - assert all(m.startswith("unet.s") for c in shared for m in c.members) - assert {c.members[0] for c in singles} == {"unet.big0", "unet.big1"} +def test_small_leaves_are_packed_with_each_other_never_welded_to_a_big_one() -> None: + """Small-FIRST is the mechanism, not a tie-break, and this is the outcome it + buys that a descending fill does not. + + Filling ascending, a leaf at or above the cell target always lands alone: + the batch before it flushes, and the next leaf is at least as large so it + overflows immediately. Filling DESCENDING, the big leaf goes down first with + room left over and the small leaves get appended to it — which welds a 4 KiB + norm's residency to a 90 MiB block, so the norm can no longer be moved + without moving 90 MiB. On this fixture that is the difference between the + twenty small leaves being one 20 MiB unit and ten of them being hostages.""" + small = [LeafCost(f"unet.s{i:02d}", MIB) for i in range(20)] + large = [LeafCost("unet.big", 90 * MIB)] + cells = pack_cells(small + large, policy=_vmm(100 * MIB), min_stream_bytes=1) + + big_cells = [c for c in cells if "unet.big" in c.members] + assert [c.members for c in big_cells] == [("unet.big",)], ( + "a leaf near the cell target must land alone, not collect small leaves" + ) + packed = [c for c in cells if "unet.big" not in c.members] + assert sum(len(c.members) for c in packed) == 20 + assert max(len(c.members) for c in packed) == 20, "and they pack together" def test_a_cell_never_mixes_forced_and_streamable_leaves() -> None: @@ -387,7 +398,7 @@ def test_a_budget_that_holds_the_hot_component_schedules_at_call_boundaries() -> assert plan.schedule is ResidencySchedule.CALL_BOUNDARY assert plan.hold_while_hot == ("unet", "text_encoder") assert plan.hot_component_bytes + plan.streams * max( - c.cast_bytes for c in plan.cells + c.cast_bytes for c in plan.cells if not c.forced ) <= plan.budget_bytes @@ -409,7 +420,7 @@ def test_the_schedule_flips_exactly_where_the_arithmetic_says() -> None: costs, budget_bytes=1 << 40, streams=2, min_stream_bytes=1, cells=_vmm(64 * MIB) ) threshold = probe.hot_component_bytes + probe.streams * max( - c.cast_bytes for c in probe.cells + c.cast_bytes for c in probe.cells if not c.forced ) at = plan_residency( costs, budget_bytes=threshold, streams=2, min_stream_bytes=1, cells=_vmm(64 * MIB) @@ -562,8 +573,10 @@ def test_the_forced_core_stays_resident_through_a_call_boundary_swap() -> None: hot = residency.hold_component("unet") assert "text_encoder.blocks.0.fc1" in hot.forced assert "text_encoder.blocks.0.fc1" not in hot.streamed - param = model.text_encoder.blocks[0].fc1.weight # type: ignore[index,union-attr] - assert not param.is_pinned() + # And it is really still where it was: a parked leaf would be sitting in + # pinned host memory instead. + params = dict(model.text_encoder.named_parameters()) + assert not params["blocks.0.fc1.weight"].is_pinned() def test_a_real_tree_reports_its_own_tax_elimination() -> None: @@ -594,3 +607,32 @@ def test_a_real_tree_reports_its_own_tax_elimination() -> None: assert b.tax_eliminated_bytes > 0 packed.release() leafwise.release() + + +def test_the_forced_core_does_not_inflate_the_call_boundary_reservation() -> None: + """Packing makes this matter. A tree with hundreds of sub-floor leaves has + ONE large forced cell, and that cell is resident in every arrangement — it + never travels. Reserving the swap ring for it would deny the call-boundary + schedule to budgets that plainly admit it.""" + costs = [LeafCost(f"unet.tiny{i:03d}", 128 * 1024) for i in range(512)] + [ + LeafCost(f"unet.l{i}", 2 * MIB) for i in range(3) + ] + # The floor forces every 128 KiB leaf resident; they pack into ONE 64 MiB + # cell, eight times larger than the only cell that can actually move. + plan = plan_residency( + costs, + budget_bytes=100 * MIB, + streams=2, + min_stream_bytes=1 * MIB, + cells=_vmm(64 * MIB), + ) + forced_cells = [c for c in plan.cells if c.forced] + movable = [c for c in plan.cells if not c.forced] + assert max(c.cast_bytes for c in forced_cells) > 8 * max( + c.cast_bytes for c in movable + ), "the packed forced core must dwarf every cell that can move" + # RED ARM: reserving the ring for the forced core would put this over budget. + assert plan.hot_component_bytes + plan.streams * max( + c.cast_bytes for c in plan.cells + ) > plan.budget_bytes + assert plan.schedule is ResidencySchedule.CALL_BOUNDARY