xiaolibai-sys
164896ed78
- nvfp4: _quantization_metadata -> comfy_quant synthesis, SlotEntry extra params (block_scale), dual-scale LoRA fold, orig_dtype fixup - LoRA: plain-mode block LoRAs no longer silently dropped; folded on GPU after weights land - loader: parallel positioned reads, pipelined group load, no per-tensor clone, header-only dtype probes - block swap: read_block intra-block parallel + disk_workers node option; async D2H write-back on transfer stream; lazy evict_all skips wasted write-back - arch: split models/loader.py out of wan_model.py; single utils/keys.py; LoRA folding consolidated in utils/lora.py; global utils/log.py ([BerniniR][Tag] format, BERNINI_LOG_LEVEL)
1565 lines
67 KiB
Python
1565 lines
67 KiB
Python
"""
|
||
Block-level GPU↔RAM offloading for Wan/Bernini transformer blocks.
|
||
|
||
Maintains a sliding window of blocks on GPU, offloading blocks outside the
|
||
window to CPU to save VRAM. A dedicated CUDA stream overlaps H2D transfers
|
||
with transformer block computation, hiding transfer latency.
|
||
|
||
When a ``RandomAccessBlockReader`` is provided (lazy-loading mode), blocks
|
||
are loaded from disk on demand by a worker thread pool. In both modes the
|
||
off-GPU blocks reside in a fixed, pre-allocated CPU "home" pool sized to
|
||
``(N - window)`` blocks, so host RAM stays flat at ``(N-W)`` for the whole
|
||
run — there is no per-step allocation churn (the previous move-based swap
|
||
produced a periodic multi-GB sawtooth on Windows).
|
||
|
||
Architecture
|
||
------------
|
||
The module is organised into five components with single responsibilities:
|
||
|
||
``_BlockWindow``
|
||
Pure state tracking — which blocks are on GPU vs CPU right now.
|
||
No CUDA calls, no tensor operations.
|
||
|
||
``_TransferEngine``
|
||
CUDA stream management + async H2D/D2H transfers. Owns the dedicated
|
||
stream, tracks in-flight prefetches, and provides safe move primitives.
|
||
|
||
``_VRAMBudget``
|
||
VRAM headroom monitoring and CUDA cache management. Decides when to
|
||
flush the allocator cache to prevent fragmentation.
|
||
|
||
``_DiskPrefetcher``
|
||
Disk→RAM loader running on a background thread pool. Reads blocks from
|
||
safetensors on demand, populates model parameters, and evicts unused
|
||
blocks from RAM.
|
||
|
||
``BlockSwapManager``
|
||
Thin orchestrator that wires the four components together and exposes
|
||
the public API used by the model's forward pass.
|
||
|
||
Usage::
|
||
|
||
mgr = BlockSwapManager(model, window_size=10, block_reader=reader)
|
||
mgr.prepare_pre_forward()
|
||
for i, block in enumerate(model.blocks):
|
||
mgr.prepare(i)
|
||
mgr.prefetch_next(i)
|
||
x = block(x, ...)
|
||
mgr.prepare_head()
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import threading
|
||
from concurrent.futures import Future, ThreadPoolExecutor
|
||
from dataclasses import dataclass, field
|
||
from typing import Optional
|
||
|
||
import torch
|
||
|
||
from .tensor_ops import (
|
||
free_module_storage,
|
||
move_module_to,
|
||
pin_module,
|
||
record_module_stream,
|
||
)
|
||
from .types import DiskLoadRequest, SlotEntry
|
||
|
||
from .log import get_logger as _get_logger
|
||
|
||
logger = _get_logger("BlockSwap")
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# _BlockWindow — pure state tracking
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
class _BlockWindow:
|
||
"""Tracks which transformer blocks are on GPU vs CPU vs disk.
|
||
|
||
Pure data structure — no CUDA calls, no tensor operations.
|
||
"""
|
||
|
||
__slots__ = ("on_gpu", "resident", "prewarmed", "total", "window_size",
|
||
"ram_resident", "disk_loading")
|
||
|
||
def __init__(self, total: int, window_size: int, prewarmed: int = 0):
|
||
self.total = total
|
||
self.window_size = max(1, window_size)
|
||
self.prewarmed = prewarmed
|
||
# Blocks known to be in CPU RAM (weights populated).
|
||
# In lazy mode this starts empty (blocks loaded on demand).
|
||
# In eager mode (legacy) this holds all blocks.
|
||
self.resident: set[int] = set(range(prewarmed, total))
|
||
# Blocks whose CPU RAM parameters have data (for eviction tracking).
|
||
# Mirrors self.resident initially; diverges in lazy mode.
|
||
self.ram_resident: set[int] = set(range(total))
|
||
# Blocks being read from disk by a background thread right now.
|
||
self.disk_loading: set[int] = set()
|
||
# Blocks currently on GPU.
|
||
self.on_gpu: set[int] = set(range(prewarmed))
|
||
|
||
# -- query ---------------------------------------------------------------
|
||
|
||
def needed(self, block_idx: int) -> set[int]:
|
||
"""Blocks that should be on GPU for the window starting at *block_idx*."""
|
||
start = max(0, block_idx)
|
||
end = min(self.total, block_idx + self.window_size)
|
||
return set(range(start, end))
|
||
|
||
def to_offload(self, block_idx: int) -> set[int]:
|
||
"""Blocks on GPU that are no longer in the window."""
|
||
return self.on_gpu - self.needed(block_idx)
|
||
|
||
def to_load(self, block_idx: int) -> set[int]:
|
||
"""Blocks needed for the window that are not yet on GPU."""
|
||
return self.needed(block_idx) - self.on_gpu
|
||
|
||
def is_on_gpu(self, block_idx: int) -> bool:
|
||
return block_idx in self.on_gpu
|
||
|
||
def is_in_ram(self, block_idx: int) -> bool:
|
||
"""True if the block's weights are loaded in CPU RAM."""
|
||
return block_idx in self.ram_resident
|
||
|
||
def prefetch_window(self, block_idx: int) -> set[int]:
|
||
"""Blocks just beyond the current window (candidate prefetch range)."""
|
||
start = block_idx + self.window_size
|
||
if start >= self.total:
|
||
return set()
|
||
count = min(self.window_size, self.total - start)
|
||
return set(range(start, start + count))
|
||
|
||
# -- mutations -----------------------------------------------------------
|
||
|
||
def mark_loaded(self, block_idx: int) -> None:
|
||
self.on_gpu.add(block_idx)
|
||
self.resident.discard(block_idx)
|
||
|
||
def mark_offloaded(self, block_idx: int) -> None:
|
||
self.on_gpu.discard(block_idx)
|
||
self.resident.add(block_idx)
|
||
|
||
def mark_ram_loaded(self, block_idx: int) -> None:
|
||
self.ram_resident.add(block_idx)
|
||
self.disk_loading.discard(block_idx)
|
||
|
||
def mark_ram_freed(self, block_idx: int) -> None:
|
||
self.ram_resident.discard(block_idx)
|
||
|
||
def mark_disk_loading(self, block_idx: int) -> None:
|
||
self.disk_loading.add(block_idx)
|
||
|
||
def clear(self) -> None:
|
||
self.on_gpu.clear()
|
||
self.resident = set(range(self.total))
|
||
self.ram_resident = set(range(self.total))
|
||
self.disk_loading.clear()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# PinStage — pinned CPU staging ring, sized to GPU ring buffer
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
@dataclass
|
||
class PinStage:
|
||
"""Pinned CPU staging buffers for truly async H2D transfers.
|
||
|
||
Each slot holds a ``dict[str, SlotEntry]`` whose tensor memory is
|
||
page-locked (pinned), so CUDA can DMA directly from it without a
|
||
staging copy. The ring cursor is shared with the GPU ring buffer:
|
||
slot ``ring_idx % ring_size`` pairs with the GPU slot at the same index.
|
||
|
||
Pinning only the transfer staging (``W + prefetch`` blocks) instead of
|
||
the entire home pool (``N - W`` blocks) keeps the pinned memory footprint
|
||
small — critical on Windows where large pinned allocations strain the
|
||
allocator.
|
||
"""
|
||
slots: list[dict | None] = field(default_factory=list)
|
||
built: list[bool] = field(default_factory=list)
|
||
|
||
@classmethod
|
||
def new(cls, ring_size: int) -> "PinStage":
|
||
return cls(slots=[None] * ring_size, built=[False] * ring_size)
|
||
|
||
@property
|
||
def size(self) -> int:
|
||
return len(self.slots)
|
||
|
||
def ensure(self, ring_idx: int, template: dict, pin_memory: bool = False) -> dict:
|
||
"""Build pinned ``SlotEntry`` for *ring_idx* from *template* (an
|
||
existing ``dict[str, SlotEntry]``, e.g. a home slot or GPU slot)
|
||
if not yet done."""
|
||
idx = ring_idx % self.size if self.size else 0
|
||
slot = self.slots[idx]
|
||
if slot is not None and self.built[idx]:
|
||
return slot
|
||
self.slots[idx] = {
|
||
n: SlotEntry.empty_like_entry(e, device="cpu", pin_memory=pin_memory)
|
||
for n, e in template.items()
|
||
}
|
||
self.built[idx] = True
|
||
return self.slots[idx]
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# _TransferEngine — CUDA stream + async H2D/D2H
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
def _copy_qt_extras(entry, qt, non_blocking: bool = False) -> None:
|
||
"""Copy a QuantizedTensor's extra layout params (e.g. nvfp4 block_scale)
|
||
into a slot entry. ``_params_extra_fields`` introspects the dataclass so
|
||
new layouts work without code changes."""
|
||
from .types import _params_extra_fields
|
||
t_extras, _ = _params_extra_fields(qt._params)
|
||
for n, t in t_extras.items():
|
||
if n in entry.extra:
|
||
entry.extra[n].copy_(t, non_blocking=non_blocking)
|
||
|
||
class _TransferEngine:
|
||
"""CUDA stream + fixed slot pools for churn-free block swaps.
|
||
|
||
Two pools are allocated **once** at construction:
|
||
|
||
* a pinned CPU "home" pool of ``n_home_slots`` block buffers — at steady
|
||
state exactly the off-GPU blocks (``total - window`` of them) occupy
|
||
these, so host RAM stays flat at ``(N-W)`` and is never returned to the
|
||
OS allocator (no periodic churn / 9 GB sawtooth);
|
||
* a GPU "slot" pool of ``n_gpu_slots`` block buffers
|
||
(``window + prefetch``) used while a block computes.
|
||
|
||
Blocks move between pools with ``copy_`` (never ``module.to``), so no new
|
||
tensors are allocated on the hot path. ``pin_memory`` enables a small
|
||
pinned staging ring (``PinStage``, sized to the prefetch count) so H2D
|
||
copies on the dedicated stream are true async DMA; the (N-W) home pool
|
||
itself stays pageable to keep the locked-page footprint small.
|
||
|
||
The legacy ``to_gpu``/``to_cpu`` *move* primitives are retained for
|
||
``StreamingBlockPipeline`` (one-off model pre-load), which does not use the
|
||
slot pools.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
device: torch.device | str,
|
||
prefetch: bool = True,
|
||
prefetch_count: int = 1,
|
||
pin_memory: bool = False,
|
||
n_home_slots: int = 0,
|
||
n_gpu_slots: int = 0,
|
||
template_block=None,
|
||
home: "BlockHome | None" = None,
|
||
):
|
||
self.device = torch.device(device) if isinstance(device, str) else device
|
||
self.prefetch = prefetch and torch.cuda.is_available()
|
||
self.prefetch_count = max(1, prefetch_count)
|
||
# pin_memory only meaningful with a CUDA device
|
||
self.pin_memory = bool(pin_memory) and torch.cuda.is_available()
|
||
self._home = home
|
||
|
||
self._stream: Optional[torch.cuda.Stream] = None
|
||
if self.prefetch:
|
||
try:
|
||
self._stream = torch.cuda.Stream(device=device)
|
||
except Exception:
|
||
pass
|
||
|
||
# block_idx → CUDA event recorded after async transfer started
|
||
self._events: dict[int, torch.cuda.Event] = {}
|
||
|
||
# hslot → CUDA event recorded after an async D2H write-back into
|
||
# that home slot completed enqueueing. Host-side consumers of the
|
||
# slot (load_block's CPU staging copy, DiskHome.read_into's disk
|
||
# write) MUST synchronise on it first; DMA-vs-DMA reuse (another
|
||
# offload into the same slot) is already ordered by the shared
|
||
# transfer stream.
|
||
self._d2h_events: dict[int, torch.cuda.Event] = {}
|
||
|
||
# ---- slot pools ----------------------------------------------------
|
||
# Each slot is a ``dict[str, SlotEntry]`` — one entry per parameter.
|
||
# ``SlotEntry`` handles both regular tensors and QuantizedTensors via
|
||
# a unified ``copy_from`` / ``assign_to`` interface, so the swap
|
||
# engine never branches on parameter type.
|
||
#
|
||
# GPU and pin (staging) slots use a shared ring buffer cursor
|
||
# (``_gpu_cursor``) — deterministic O(1) allocation.
|
||
nh = max(0, n_home_slots)
|
||
ng = max(0, n_gpu_slots)
|
||
self._home_pool: list[dict | None] = [None] * nh
|
||
self._home_built: list[bool] = [False] * nh
|
||
self._home_free: list[int] = list(range(nh))
|
||
|
||
self._pin: PinStage = PinStage.new(max(1, self.prefetch_count)) # pinned CPU staging (small — only for prefetch overlap)
|
||
self._gpu_pool: list[dict | None] = [None] * ng # GPU ring buffer
|
||
self._gpu_built: list[bool] = [False] * ng
|
||
|
||
self._gpu_cursor: int = 0
|
||
self._block_home: dict[int, int] = {}
|
||
self._block_gpu: dict[int, int] = {}
|
||
|
||
# -- slot bookkeeping ----------------------------------------------------
|
||
|
||
def _ring_acquire(self) -> int:
|
||
"""Advance the GPU ring-buffer cursor and return the next slot index."""
|
||
ng = len(self._gpu_pool)
|
||
idx = self._gpu_cursor % ng
|
||
self._gpu_cursor = (self._gpu_cursor + 1) % ng
|
||
return idx
|
||
|
||
def _ensure_entries(self, pool, built, idx, block, device) -> dict:
|
||
"""Build a ``name → SlotEntry`` dict for *block* if not already done.
|
||
|
||
On the first build, any LoRA payload co-located on the block
|
||
(``block._lora_payload``, set by the streaming/disk loader) is folded
|
||
into the slot and the payload is cleared — so the slot itself becomes
|
||
the single owner of that block's LoRA, and rebuilds never re-attach it.
|
||
"""
|
||
slot = pool[idx]
|
||
if slot is not None and built[idx]:
|
||
return slot
|
||
entries = {
|
||
n: SlotEntry.empty_like(
|
||
p, device=device,
|
||
# Home pool stays pageable: only the small PinStage
|
||
# staging ring is pinned. Pinning the whole (N-W) home
|
||
# pool can fail outright (cudaHostAlloc OOM) on machines
|
||
# whose RAM is mostly occupied by the resident model.
|
||
pin_memory=False,
|
||
)
|
||
for n, p in block.named_parameters()
|
||
}
|
||
payload = getattr(block, "_lora_payload", None)
|
||
if payload:
|
||
for pname, lora_list in payload.items():
|
||
if pname in entries:
|
||
entries[pname].lora = lora_list
|
||
block._lora_payload = None # consumed into the slot
|
||
pool[idx] = entries
|
||
built[idx] = True
|
||
return pool[idx]
|
||
|
||
def _swap_param(self, block, entries, non_blocking: bool = False) -> None:
|
||
"""Copy every param's data from *block* into *entries*, then assign
|
||
the entries back to the block's ``_parameters`` — all via ``SlotEntry``,
|
||
which handles both regular and QuantizedTensor params uniformly."""
|
||
for n, p in block.named_parameters():
|
||
e = entries.get(n)
|
||
if e is None:
|
||
continue
|
||
e.copy_from(p, non_blocking=non_blocking)
|
||
e.assign_to(block, n)
|
||
|
||
def reset(self) -> None:
|
||
"""Return every slot to the free list and clear block mappings."""
|
||
self._block_home.clear()
|
||
self._block_gpu.clear()
|
||
self._home_free = list(range(len(self._home_pool)))
|
||
self._gpu_cursor = 0
|
||
self._events.clear()
|
||
self._d2h_events.clear()
|
||
|
||
def clear_pools(self) -> None:
|
||
"""Drop all pool tensor references so GC can reclaim RAM/VRAM."""
|
||
nh, ng = len(self._home_pool), len(self._gpu_pool)
|
||
self._home_pool = [None] * nh
|
||
self._gpu_pool = [None] * ng
|
||
self._pin = PinStage.new(ng)
|
||
self._home_built = [False] * nh
|
||
self._gpu_built = [False] * ng
|
||
self._block_home.clear()
|
||
self._block_gpu.clear()
|
||
self._events.clear()
|
||
self._d2h_events.clear()
|
||
|
||
def _wait_home_slot(self, hslot: int) -> None:
|
||
"""Host-wait for any in-flight async D2H into home slot *hslot*.
|
||
|
||
Required before the slot's tensors are read or overwritten by the
|
||
CPU (load staging copy, disk read_into) — a CUDA stream provides no
|
||
ordering guarantee against host-side memory access.
|
||
"""
|
||
ev = self._d2h_events.pop(hslot, None)
|
||
if ev is not None:
|
||
ev.synchronize()
|
||
|
||
def forget_home(self, block_idx: int) -> None:
|
||
"""Drop the home-slot mapping for *block_idx*."""
|
||
self._block_home.pop(block_idx, None)
|
||
|
||
def discard_block(self, block_idx: int) -> None:
|
||
"""Drop *block_idx* from the GPU ring WITHOUT writing weights home.
|
||
|
||
Safe only when the weights are recoverable elsewhere (lazy disk
|
||
mode) and the pools are about to be cleared — the block's params are
|
||
left pointing at GPU ring slots that ``clear_pools`` releases.
|
||
Steady-state offloads must NOT use this: the recycled home slot's
|
||
old content is dead, so the write-back is the only off-GPU copy.
|
||
"""
|
||
self._events.pop(block_idx, None)
|
||
self._block_gpu.pop(block_idx, None)
|
||
|
||
def _acquire_home(self) -> int:
|
||
if not self._home_free:
|
||
raise RuntimeError("block_swap: out of home slots")
|
||
return self._home_free.pop()
|
||
|
||
def init_gpu(self, block_idx: int, block) -> None:
|
||
"""Place *block*'s weights onto a GPU ring slot at construction time."""
|
||
gslot = self._ring_acquire()
|
||
entries = self._ensure_entries(self._gpu_pool, self._gpu_built,
|
||
gslot, block, self.device)
|
||
self._swap_param(block, entries)
|
||
self._block_gpu[block_idx] = gslot
|
||
|
||
def init_home(self, block_idx: int, block) -> None:
|
||
"""Place *block*'s weights onto a home slot at construction time."""
|
||
hslot = self._acquire_home()
|
||
entries = self._ensure_entries(self._home_pool, self._home_built,
|
||
hslot, block, "cpu")
|
||
self._swap_param(block, entries)
|
||
self._block_home[block_idx] = hslot
|
||
|
||
# -- legacy move primitives (StreamingBlockPipeline) --------------------
|
||
|
||
def to_gpu(self, module: torch.nn.Module, non_blocking: bool = False) -> None:
|
||
"""Move *module* to GPU, optionally async on the transfer stream."""
|
||
if non_blocking and self._stream is not None:
|
||
self._stream.wait_stream(torch.cuda.current_stream())
|
||
if self.pin_memory:
|
||
pin_module(module)
|
||
with torch.cuda.stream(self._stream):
|
||
move_module_to(module, self.device, non_blocking=True)
|
||
record_module_stream(module, self._stream)
|
||
else:
|
||
move_module_to(module, self.device)
|
||
|
||
def to_cpu(
|
||
self,
|
||
module: torch.nn.Module,
|
||
block_idx: int | None = None,
|
||
) -> None:
|
||
"""Move *module* to CPU, waiting for its own in-flight prefetch if needed.
|
||
|
||
Only the block being offloaded is waited on. Other blocks' in-flight
|
||
prefetches are intentionally left alone — they are still useful and will
|
||
be consumed by ``sync_prefetch`` when those blocks enter the window.
|
||
Cancelling them here (the previous behaviour) silently disabled all
|
||
prefetch overlap and forced every load back to synchronous.
|
||
"""
|
||
if block_idx is not None and block_idx in self._events:
|
||
event = self._events.pop(block_idx)
|
||
torch.cuda.current_stream().wait_event(event)
|
||
move_module_to(module, "cpu")
|
||
|
||
# -- slot-based block transfer (churn-free) -----------------------------
|
||
|
||
def ensure_home(self, block_idx: int, block) -> int:
|
||
"""Ensure *block_idx* has a populated home slot; return its index."""
|
||
if block_idx in self._block_home:
|
||
return self._block_home[block_idx]
|
||
hslot = self._acquire_home()
|
||
# The acquired slot may still be receiving an async D2H from a
|
||
# previous offload — DiskHome.read_into overwrites it from the host.
|
||
self._wait_home_slot(hslot)
|
||
entries = self._ensure_entries(self._home_pool, self._home_built,
|
||
hslot, block, "cpu")
|
||
if self._home is not None:
|
||
self._home.read_into(block_idx, entries, block)
|
||
self._block_home[block_idx] = hslot
|
||
return hslot
|
||
|
||
def load_block(self, block_idx: int, block, non_blocking: bool = True) -> None:
|
||
"""Copy *block*'s weights onto a GPU ring slot.
|
||
|
||
When ``pin_memory`` is enabled, data flows through a pinned CPU
|
||
staging ring for truly async H2D copies. Otherwise it copies
|
||
directly from block parameters to the GPU ring.
|
||
"""
|
||
if block_idx in self._block_gpu:
|
||
return
|
||
hslot = self.ensure_home(block_idx, block)
|
||
# Block params point at the home slot (assigned at offload); the CPU
|
||
# staging copies below read them on the host, so any in-flight async
|
||
# D2H into that slot must finish first.
|
||
self._wait_home_slot(hslot)
|
||
ridx = self._ring_acquire()
|
||
|
||
gpu = self._ensure_entries(self._gpu_pool, self._gpu_built,
|
||
ridx, block, self.device)
|
||
nb = self._stream is not None and non_blocking
|
||
|
||
if self.pin_memory and non_blocking and self._stream is not None:
|
||
# ── Pinned staging path (prefetch only) ──────────────────────
|
||
pin = self._pin.ensure(ridx, gpu, pin_memory=True)
|
||
# Step 1: block params → pinned CPU staging
|
||
for n, p in block.named_parameters():
|
||
e = pin.get(n)
|
||
if e is not None:
|
||
e.copy_from(p)
|
||
# Step 2: pinned CPU → GPU on transfer stream
|
||
self._stream.wait_stream(torch.cuda.current_stream())
|
||
with torch.cuda.stream(self._stream):
|
||
for n in pin:
|
||
pe, ge = pin[n], gpu.get(n)
|
||
if ge is None:
|
||
continue
|
||
if pe.is_qt:
|
||
ge.data.copy_(pe.data, non_blocking=True)
|
||
ge.scale.copy_(pe.scale, non_blocking=True)
|
||
for en, et in pe.extra.items():
|
||
if en in ge.extra:
|
||
ge.extra[en].copy_(et, non_blocking=True)
|
||
else:
|
||
ge.data.copy_(pe.data, non_blocking=True)
|
||
# Sync: compute stream waits for transfer stream BEFORE assign_to
|
||
compute = torch.cuda.current_stream()
|
||
ev = torch.cuda.Event()
|
||
ev.record(self._stream)
|
||
compute.wait_event(ev)
|
||
nb = False # pin path already synced; skip outer event recording
|
||
elif self.pin_memory and non_blocking:
|
||
# Pin enabled but no stream — fallback to direct sync copy
|
||
for n, p in block.named_parameters():
|
||
e = gpu.get(n)
|
||
if e is None:
|
||
continue
|
||
if e.is_qt:
|
||
e.data.copy_(p.data._qdata)
|
||
e.scale.copy_(p.data._params.scale)
|
||
_copy_qt_extras(e, p.data, non_blocking=False)
|
||
else:
|
||
e.data.copy_(p.data)
|
||
# No event needed — synchronous copy completed above.
|
||
nb = False
|
||
else:
|
||
# ── Direct path (no pin pool) ────────────────────────────────
|
||
for n, p in block.named_parameters():
|
||
e = gpu.get(n)
|
||
if e is None:
|
||
continue
|
||
if e.is_qt:
|
||
e.data.copy_(p.data._qdata, non_blocking=nb)
|
||
e.scale.copy_(p.data._params.scale, non_blocking=nb)
|
||
_copy_qt_extras(e, p.data, non_blocking=nb)
|
||
else:
|
||
e.data.copy_(p.data, non_blocking=nb)
|
||
|
||
# Assign GPU entries → block _parameters
|
||
for n in gpu:
|
||
gpu[n].assign_to(block, n)
|
||
|
||
# Carry the co-located LoRA payload along with the block weights
|
||
# (unified slot: home and GPU copies both hold it until folded).
|
||
hslot = self._block_home.get(block_idx)
|
||
if hslot is not None:
|
||
home_entries = self._home_pool[hslot]
|
||
for n in gpu:
|
||
he = home_entries.get(n)
|
||
if he is not None and he.lora:
|
||
gpu[n].lora = he.lora
|
||
|
||
# Recycle home slot
|
||
hslot = self._block_home.pop(block_idx, None)
|
||
if hslot is not None:
|
||
self._home_free.append(hslot)
|
||
self._block_gpu[block_idx] = ridx
|
||
if nb:
|
||
event = torch.cuda.Event()
|
||
event.record(self._stream)
|
||
self._events[block_idx] = event
|
||
|
||
def offload_block(self, block_idx: int, block, force: bool = False) -> bool:
|
||
"""Copy *block* back to a home slot; returns True if moved.
|
||
|
||
Copies GPU slot entry → home slot entry directly (NOT through
|
||
``block.named_parameters()``, which may have been freed by
|
||
``free_module_storage`` during model switching).
|
||
|
||
The D2H copies are issued on the dedicated transfer stream
|
||
(``non_blocking=True``) so they overlap with subsequent compute
|
||
instead of stalling the compute stream with ~one-block-worth of
|
||
synchronous pageable copies per offload. Completion is tracked per
|
||
home slot in ``_d2h_events``; host-side consumers synchronise via
|
||
``_wait_home_slot``, while DMA-side reuse is ordered by the shared
|
||
transfer stream itself. ``assign_to`` below is a host-side pointer
|
||
rebind only, so it is safe to run while the copy is in flight.
|
||
"""
|
||
event = self._events.pop(block_idx, None)
|
||
if event is not None:
|
||
torch.cuda.current_stream().wait_event(event)
|
||
gslot = self._block_gpu.pop(block_idx, None)
|
||
if gslot is None:
|
||
return False
|
||
hslot = self._block_home.get(block_idx)
|
||
if hslot is None:
|
||
if self._home_free:
|
||
hslot = self._home_free.pop()
|
||
elif force:
|
||
victim = min(self._block_home.keys())
|
||
hslot = self._block_home.pop(victim)
|
||
# A recycled slot may still be receiving the victim's async
|
||
# D2H; the force path releases + reassigns it from the host.
|
||
self._wait_home_slot(hslot)
|
||
if self._home is not None:
|
||
self._home.release(victim)
|
||
else:
|
||
self._block_gpu[block_idx] = gslot
|
||
return False
|
||
# Build (or reuse) home slot entries from the block — at this point
|
||
# the block may already have been freed (empty(0) params), so we
|
||
# build from the GPU slot's metadata instead.
|
||
gpu_entries = self._gpu_pool[gslot]
|
||
home_entries = self._ensure_entries(self._home_pool, self._home_built,
|
||
hslot, block, "cpu")
|
||
|
||
if self._stream is not None:
|
||
# ── Async D2H on the transfer stream ─────────────────────────
|
||
# The transfer stream must first drain the compute stream (the
|
||
# offloaded block's forward may still be reading its GPU slot).
|
||
self._stream.wait_stream(torch.cuda.current_stream())
|
||
with torch.cuda.stream(self._stream):
|
||
for n in list(gpu_entries.keys() & home_entries.keys()):
|
||
ge, he = gpu_entries[n], home_entries[n]
|
||
if ge.is_qt:
|
||
he.data.copy_(ge.data, non_blocking=True)
|
||
he.scale.copy_(ge.scale, non_blocking=True)
|
||
for en, et in ge.extra.items():
|
||
if en in he.extra:
|
||
he.extra[en].copy_(et, non_blocking=True)
|
||
else:
|
||
he.data.copy_(ge.data, non_blocking=True)
|
||
he.lora = ge.lora # carry co-located LoRA (None once folded)
|
||
he.assign_to(block, n)
|
||
ev = torch.cuda.Event()
|
||
ev.record(self._stream)
|
||
self._d2h_events[hslot] = ev
|
||
else:
|
||
# ── Synchronous fallback (no transfer stream) ────────────────
|
||
for n in list(gpu_entries.keys() & home_entries.keys()):
|
||
ge, he = gpu_entries[n], home_entries[n]
|
||
if ge.is_qt:
|
||
he.data.copy_(ge.data)
|
||
he.scale.copy_(ge.scale)
|
||
for en, et in ge.extra.items():
|
||
if en in he.extra:
|
||
he.extra[en].copy_(et)
|
||
else:
|
||
he.data.copy_(ge.data)
|
||
he.lora = ge.lora # carry co-located LoRA (None once folded)
|
||
he.assign_to(block, n)
|
||
self._block_home[block_idx] = hslot
|
||
return True
|
||
|
||
# -- prefetch ------------------------------------------------------------
|
||
|
||
def start_prefetch(self, block_idx: int, block) -> None:
|
||
"""Async-load *block* ahead of the window (no-op if already loaded)."""
|
||
if block_idx in self._block_gpu or block_idx in self._events:
|
||
return
|
||
self.load_block(block_idx, block, non_blocking=True)
|
||
|
||
def sync_prefetch(self, block_idx: int) -> bool:
|
||
"""Wait for *block_idx*'s async transfer to finish. Returns True if waited."""
|
||
event = self._events.pop(block_idx, None)
|
||
if event is None:
|
||
return False
|
||
torch.cuda.current_stream().wait_event(event)
|
||
return True
|
||
|
||
def cancel_all(self) -> None:
|
||
"""Abort all in-flight prefetches."""
|
||
if not self._events:
|
||
return
|
||
if self._stream is not None:
|
||
self._stream.synchronize()
|
||
self._events.clear()
|
||
|
||
def sync_all(self) -> None:
|
||
"""Block until every CUDA stream (including compute) is done."""
|
||
if not torch.cuda.is_available():
|
||
return
|
||
try:
|
||
torch.cuda.synchronize()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# _VRAMBudget — memory monitoring
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
class _VRAMBudget:
|
||
"""Tracks per-block VRAM usage and manages CUDA cache flushes.
|
||
|
||
Decides when to call ``torch.cuda.empty_cache()`` to defragment the
|
||
allocator — too often hurts perf, too rarely causes OOM from
|
||
fragmentation holes.
|
||
"""
|
||
|
||
def __init__(self, block_mb: float):
|
||
self.block_mb = block_mb
|
||
|
||
@staticmethod
|
||
def estimate_block_mb(block) -> float:
|
||
"""Rough VRAM for one block in MB (assumes bf16/fp16 = 2 bytes/param)."""
|
||
total = sum(p.numel() for p in block.parameters()) * 2
|
||
return total / (1024 * 1024)
|
||
|
||
def free_mb(self, device: torch.device | str | None = None) -> float:
|
||
if not torch.cuda.is_available():
|
||
return float("inf")
|
||
try:
|
||
free, _ = torch.cuda.mem_get_info(device)
|
||
return free / (1024 * 1024)
|
||
except Exception:
|
||
return float("inf")
|
||
|
||
def is_fragmented(self, device: torch.device | str | None = None) -> bool:
|
||
if not torch.cuda.is_available():
|
||
return False
|
||
try:
|
||
allocated = torch.cuda.memory_allocated(device) / (1024 * 1024)
|
||
reserved = torch.cuda.memory_reserved(device) / (1024 * 1024)
|
||
return reserved > allocated + self.block_mb * 2
|
||
except Exception:
|
||
return False
|
||
|
||
def maybe_flush(
|
||
self,
|
||
device: torch.device | str | None = None,
|
||
reserve_blocks: int = 2,
|
||
) -> None:
|
||
"""Release cached GPU memory if free space is tight or fragmented."""
|
||
if not torch.cuda.is_available():
|
||
return
|
||
needed_mb = self.block_mb * max(1, reserve_blocks)
|
||
if self.free_mb(device) < needed_mb or self.is_fragmented(device):
|
||
try:
|
||
torch.cuda.empty_cache()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# _DiskPrefetcher — lazy disk→RAM→GPU pipeline
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
class _DiskPrefetcher:
|
||
"""Background disk→RAM loader with on-demand CPU RAM eviction.
|
||
|
||
Reads block weights from safetensors on a thread pool, populates the
|
||
model's parameter data, and frees unused blocks' CPU RAM when they
|
||
leave the GPU window + prefetch range.
|
||
|
||
Parameters
|
||
----------
|
||
model:
|
||
The diffusion model (must have ``.blocks`` list and ``_lora_groups``).
|
||
block_reader:
|
||
Random-access safetensors reader.
|
||
max_workers:
|
||
Thread pool size (typically 2).
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model,
|
||
block_reader,
|
||
lora_reader=None,
|
||
max_workers: int = 2,
|
||
):
|
||
self._model = model
|
||
self._blocks = model.blocks
|
||
self._reader = block_reader
|
||
self._lora_reader = lora_reader
|
||
|
||
self._executor = ThreadPoolExecutor(
|
||
max_workers=max_workers,
|
||
thread_name_prefix="disk-prefetch",
|
||
)
|
||
self._pending: dict[int, Future] = {} # block_idx → Future
|
||
# Blocks whose weights are confirmed resident in CPU RAM (worker
|
||
# finished, or a synchronous load completed). Lets ``ensure_ram``
|
||
# return immediately instead of re-reading from disk when a completed
|
||
# future was reclaimed from ``_pending``.
|
||
self._loaded: set[int] = set()
|
||
self._lock = threading.Lock()
|
||
self._shut_down = False
|
||
|
||
def ensure_ram(self, block_idx: int) -> None:
|
||
"""Guarantee block *block_idx* is in CPU RAM (block until ready)."""
|
||
# Already resident? No I/O, no locking past the check.
|
||
with self._lock:
|
||
if block_idx in self._loaded:
|
||
return
|
||
fut = self._pending.pop(block_idx, None)
|
||
if fut is not None:
|
||
fut.result() # wait for disk read (re-raises worker errors)
|
||
else:
|
||
# Not loading yet — load synchronously on this thread.
|
||
self._load_immediate(
|
||
DiskLoadRequest(block_idx=block_idx, group_key=f"blocks.{block_idx}")
|
||
)
|
||
with self._lock:
|
||
self._loaded.add(block_idx)
|
||
|
||
def start_ram_load(self, block_idx: int) -> None:
|
||
"""Begin loading *block_idx* from disk in a background thread."""
|
||
if self._shut_down:
|
||
return # pool torn down; caller must not dispatch new work
|
||
req = DiskLoadRequest(block_idx=block_idx, group_key=f"blocks.{block_idx}")
|
||
with self._lock:
|
||
# Skip if already resident or an in-flight load exists. Done
|
||
# futures stay in ``_pending`` until ``ensure_ram`` consumes them,
|
||
# so a completed prefetch is never silently dropped and re-read.
|
||
if block_idx in self._loaded or block_idx in self._pending:
|
||
return
|
||
fut = self._executor.submit(self._load_immediate, req)
|
||
self._pending[block_idx] = fut
|
||
fut.add_done_callback(self._on_load_done)
|
||
|
||
def evict_ram(self, block_idx: int) -> None:
|
||
"""Free block *block_idx*'s CPU RAM, discarding its weight data.
|
||
|
||
Safe to call only after the block's GPU copy is confirmed.
|
||
"""
|
||
try:
|
||
block = self._blocks[block_idx]
|
||
except IndexError:
|
||
return
|
||
free_module_storage(block)
|
||
|
||
def evict_all_ram(self) -> None:
|
||
"""Free ALL blocks' CPU RAM."""
|
||
for block_idx in range(len(self._blocks)):
|
||
self.evict_ram(block_idx)
|
||
with self._lock:
|
||
self._loaded.clear()
|
||
|
||
def cancel_all(self) -> None:
|
||
"""Abort all in-flight disk reads."""
|
||
with self._lock:
|
||
for fut in self._pending.values():
|
||
fut.cancel()
|
||
self._pending.clear()
|
||
|
||
def _on_load_done(self, fut: Future) -> None:
|
||
"""Record a successfully completed load so ``ensure_ram`` won't re-read.
|
||
|
||
Bound to each future via ``add_done_callback``; runs on the worker
|
||
thread. Errors are deliberately *not* swallowed here — they surface
|
||
when ``ensure_ram`` calls ``fut.result()`` (or trigger a synchronous
|
||
re-read that re-raises).
|
||
"""
|
||
if fut.exception() is None:
|
||
with self._lock:
|
||
# ``_pending`` may already hold a newer entry for this block;
|
||
# ``_loaded`` is keyed by block_idx, so just mark it resident.
|
||
for bi in list(self._pending):
|
||
if self._pending.get(bi) is fut:
|
||
self._loaded.add(bi)
|
||
break
|
||
|
||
def join(self, timeout: float = 30.0) -> None:
|
||
"""Block until every in-flight disk load finishes (no pool teardown).
|
||
|
||
Call *before* freeing block RAM: a worker still executing
|
||
``_load_immediate`` could otherwise write into a parameter that
|
||
``evict_all_ram`` has shrunk to ``empty(0)`` -> copy_ crash or NaN.
|
||
"""
|
||
with self._lock:
|
||
futs = list(self._pending.values())
|
||
for f in futs:
|
||
try:
|
||
f.result(timeout=timeout)
|
||
except Exception:
|
||
# Worker error — let ensure_ram / the forward re-raise if the
|
||
# block is actually needed; nothing to free here.
|
||
pass
|
||
|
||
def shutdown(self) -> None:
|
||
"""Cancel I/O and shut down the thread pool.
|
||
|
||
``wait=True`` guarantees every worker has finished ``_load_immediate``
|
||
(and its ``_stream_load_group`` writes into model params) *before* the
|
||
caller frees any block storage. ``Future.cancel()`` is a no-op for
|
||
already-running tasks, so without this a worker could still be writing
|
||
while ``evict_all_ram()`` swaps ``param.data`` for an ``empty(0)``
|
||
tensor -> copy_ crash or NaN weights.
|
||
"""
|
||
try:
|
||
self.cancel_all()
|
||
try:
|
||
self._executor.shutdown(wait=True, timeout=30)
|
||
except Exception:
|
||
# Last-resort: don't block teardown if the pool is wedged.
|
||
self._executor.shutdown(wait=False)
|
||
# Belt-and-suspenders: ensures any still-running worker finished
|
||
# writing before evict_all_ram() frees the parameters.
|
||
self.join()
|
||
self.evict_all_ram()
|
||
finally:
|
||
# Always mark the prefetcher dead so evict_all()/start_ram_load()
|
||
# can't dispatch or free against a torn-down pool.
|
||
self._shut_down = True
|
||
|
||
# ── internal ───────────────────────────────────────────────────────
|
||
|
||
def _load_immediate(self, req: DiskLoadRequest) -> None:
|
||
"""Read *req.block_idx* from disk into CPU RAM (called from thread)."""
|
||
from ..models.loader import _stream_load_group
|
||
try:
|
||
tensors = self._reader.read_block(req.block_idx)
|
||
except Exception as e:
|
||
logger.error(
|
||
"Failed to read block %d: %s", req.block_idx, e
|
||
)
|
||
raise
|
||
# Fused LoRA approach: skip fold at disk-load time;
|
||
# BlockSwapManager folds LoRA after the block lands on the GPU ring
|
||
# buffer. Co-locate this block's LoRA with its slot so the unified
|
||
# (block, lora) → (block) merge happens on the GPU ring, and the
|
||
# global LoRA pool can be dropped once every block has been attached.
|
||
_stream_load_group(
|
||
self._model, req.group_key, tensors, None
|
||
)
|
||
self._attach_block_lora(req.block_idx)
|
||
|
||
def _attach_block_lora(self, block_idx: int) -> None:
|
||
"""Co-locate *block_idx*'s LoRA into ``block._lora_payload``.
|
||
|
||
Reads the block's LoRA on demand from the ``LoraBlockReader`` (lazy
|
||
mode) so the global LoRA RAM pool is never built. The payload is
|
||
consumed once by ``_TransferEngine._ensure_entries`` when the slot is
|
||
first built.
|
||
"""
|
||
reader = self._lora_reader
|
||
if reader is None:
|
||
return
|
||
try:
|
||
payload = reader.read_block_lora(block_idx)
|
||
except Exception:
|
||
logger.error("Failed to read LoRA for block %d", block_idx)
|
||
raise
|
||
if payload:
|
||
try:
|
||
self._blocks[block_idx]._lora_payload = payload
|
||
except IndexError:
|
||
pass
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# BlockHome — unified block-weight source (resident RAM or disk)
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
class BlockHome:
|
||
"""Source of a transformer block's weights.
|
||
|
||
Resident mode (weights already in CPU RAM) and lazy mode (weights on
|
||
disk) are expressed through this single interface, so the swap engine
|
||
never branches on mode. The engine owns a fixed pool of CPU "home"
|
||
buffers; ``read_into`` populates one such buffer with a block's weights
|
||
on demand (resident: already there; disk: stream from safetensors).
|
||
"""
|
||
|
||
def __init__(self, model=None):
|
||
self._model = model
|
||
|
||
def read_into(self, block_idx: int, slot: dict, model_blocks) -> None:
|
||
"""Fill *slot* (``name -> Tensor``) with *block_idx*'s weights."""
|
||
raise NotImplementedError
|
||
|
||
def release(self, block_idx: int) -> None:
|
||
"""Drop cached weights for *block_idx* (no-op for resident RAM)."""
|
||
pass
|
||
|
||
def shutdown(self) -> None:
|
||
"""Tear down any background I/O (no-op for resident RAM)."""
|
||
pass
|
||
|
||
def join(self) -> None:
|
||
"""Wait for in-flight I/O (no-op for resident RAM)."""
|
||
pass
|
||
|
||
|
||
class RamHome(BlockHome):
|
||
"""Resident mode: weights live in the model's CPU parameters.
|
||
|
||
At manager init every block's weights are copied into a fixed home
|
||
buffer owned by the transfer engine; afterwards the weights always
|
||
reside in that buffer, so ``read_into`` is a no-op.
|
||
"""
|
||
|
||
def read_into(self, block_idx: int, slot: dict, model_blocks) -> None:
|
||
# Weights already live in the engine's home buffer (populated at
|
||
# init, or written back on offload). Nothing to fetch.
|
||
return
|
||
|
||
|
||
class DiskHome(BlockHome):
|
||
"""Lazy mode: weights are streamed from safetensors on demand."""
|
||
|
||
def __init__(self, model, block_reader, lora_reader=None, max_workers: int = 4):
|
||
self._model = model
|
||
self._disk = _DiskPrefetcher(
|
||
model, block_reader, lora_reader=lora_reader, max_workers=max_workers,
|
||
)
|
||
|
||
def read_into(self, block_idx: int, slot: dict, block) -> None:
|
||
# *block* is the manager's model block (identical object to
|
||
# ``self._model.blocks[block_idx]``); ``ensure_ram`` populates its
|
||
# parameters in place from safetensors.
|
||
#
|
||
# The home slot's ``SlotEntry`` entries are pre-allocated by
|
||
# ``_ensure_entries`` from the block's initial (possibly
|
||
# pre-materialised) parameter set. For fp8 deferred-weight ops the
|
||
# block's ``named_parameters()`` is empty *before* ``ensure_ram``
|
||
# and complete *afterwards* — ``SlotEntry.copy_from`` handles both
|
||
# regular and quantized tensor types transparently.
|
||
if block_idx not in self._disk._loaded:
|
||
self._disk.ensure_ram(block_idx)
|
||
for n, p in block.named_parameters():
|
||
e = slot.get(n)
|
||
if e is None:
|
||
# Deferred-weight materialisation: the param didn't exist
|
||
# when the slot was built. Build a new CPU SlotEntry now.
|
||
# Home pool stays pageable (see _ensure_entries); only the
|
||
# PinStage staging ring is pinned.
|
||
slot[n] = SlotEntry.empty_like(
|
||
p, device="cpu", pin_memory=False,
|
||
)
|
||
slot[n].copy_from(p)
|
||
else:
|
||
e.copy_from(p)
|
||
|
||
def release(self, block_idx: int) -> None:
|
||
self._disk.evict_ram(block_idx)
|
||
self._disk._loaded.discard(block_idx)
|
||
|
||
def shutdown(self) -> None:
|
||
self._disk.shutdown()
|
||
|
||
def join(self) -> None:
|
||
self._disk.join()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# BlockSwapManager — public API
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
class BlockSwapManager:
|
||
"""Sliding-window GPU block cache with async CUDA-stream prefetch.
|
||
|
||
Keeps at most *window_size* consecutive transformer blocks on GPU at any
|
||
time. Non-block modules are brought to GPU on demand. A dedicated CUDA
|
||
stream prefetches upcoming blocks while the current block computes.
|
||
|
||
Parameters
|
||
----------
|
||
model:
|
||
The diffusion model (must have ``.blocks`` list and peripheral attrs).
|
||
window_size:
|
||
Number of transformer blocks to keep on GPU simultaneously.
|
||
prefetch:
|
||
If True, async H2D copies run on a dedicated CUDA stream.
|
||
prefetch_count:
|
||
Max blocks to prefetch ahead of the current window.
|
||
pin_memory:
|
||
If True, pin CPU copies before async transfer (uses more host RAM).
|
||
"""
|
||
|
||
# Peripheral module names managed alongside transformer blocks.
|
||
_PERIPHERAL_NAMES = (
|
||
"patch_embedding",
|
||
"text_embedding",
|
||
"time_embedding",
|
||
"time_projection",
|
||
"head",
|
||
)
|
||
|
||
def __init__(
|
||
self,
|
||
model,
|
||
window_size: int = 10,
|
||
device: str = "cuda",
|
||
prefetch: bool = True,
|
||
prefetch_count: int = 1,
|
||
pin_memory: bool = True,
|
||
block_reader=None,
|
||
lora_reader=None,
|
||
max_disk_workers: int = 4,
|
||
):
|
||
self.model = model
|
||
self._blocks = model.blocks
|
||
self._prefetch = bool(prefetch)
|
||
total = len(self._blocks)
|
||
|
||
# Resolve pre-warmed count. Resident mode (no block_reader) pre-warms
|
||
# a full window onto GPU so that, at steady state, exactly
|
||
# ``total - window`` blocks rest in the CPU home pool — host RAM stays
|
||
# flat at ``(N-W)`` with no periodic allocation churn. Lazy mode
|
||
# pre-warms nothing (blocks are read from disk on demand).
|
||
prewarmed = window_size if block_reader is None else 0
|
||
|
||
# Unified weight source: RamHome (resident RAM) or DiskHome (lazy
|
||
# disk). The swap engine never branches on mode afterwards.
|
||
if block_reader is not None:
|
||
self._home: BlockHome = DiskHome(
|
||
model, block_reader, lora_reader=lora_reader,
|
||
max_workers=max_disk_workers,
|
||
)
|
||
self._lazy = True
|
||
else:
|
||
self._home: BlockHome = RamHome(model)
|
||
self._lazy = False
|
||
|
||
# ---- sub-components ------------------------------------------------
|
||
self._window = _BlockWindow(total, window_size, prewarmed)
|
||
self._xfer = _TransferEngine(
|
||
device, prefetch, prefetch_count, pin_memory,
|
||
n_home_slots=max(total - window_size, 1),
|
||
n_gpu_slots=window_size + prefetch_count,
|
||
template_block=self._blocks[0] if total else None,
|
||
home=self._home,
|
||
)
|
||
|
||
# Per-block VRAM estimate.
|
||
_meta = getattr(model, "_block_meta", None)
|
||
if _meta and _meta.get("avg_mb", 0) > 0:
|
||
block_mb = _meta["avg_mb"]
|
||
else:
|
||
block_mb = _VRAMBudget.estimate_block_mb(self._blocks[0])
|
||
self._budget = _VRAMBudget(block_mb)
|
||
|
||
# Peripheral modules (embeddings, head, etc.)
|
||
self._peri = {
|
||
name: getattr(model, name, None)
|
||
for name in self._PERIPHERAL_NAMES
|
||
}
|
||
|
||
logger.info(
|
||
"%d blocks total, window=%d, per-block ~%.0f MB, "
|
||
"target VRAM ~%.0f MB, prefetch=%s, count=%d, pin=%s, "
|
||
"prewarmed=%d, lazy=%s",
|
||
total, window_size, block_mb, window_size * block_mb,
|
||
prefetch, prefetch_count, pin_memory, prewarmed, self._lazy,
|
||
)
|
||
|
||
self._init_state()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Public attribute accessors (used by the model's forward pass)
|
||
# ------------------------------------------------------------------
|
||
|
||
@property
|
||
def window(self) -> int:
|
||
"""Sliding-window size in blocks (GPU-resident count).
|
||
|
||
Mirrors ``_window.window_size`` so the forward pass can cheaply
|
||
check whether the configured window changed without reaching into
|
||
the internal ``_window`` state object.
|
||
"""
|
||
return self._window.window_size
|
||
|
||
@property
|
||
def device(self) -> "torch.device | str":
|
||
"""GPU compute device blocks/activations are streamed to.
|
||
|
||
Mirrors ``_xfer.device`` (defaults to ``"cuda"`` when the manager
|
||
is constructed without an explicit device).
|
||
"""
|
||
return self._xfer.device
|
||
|
||
# ------------------------------------------------------------------
|
||
# Initialisation
|
||
# ------------------------------------------------------------------
|
||
|
||
def _init_state(self) -> None:
|
||
"""Establish deterministic starting state.
|
||
|
||
Reset the slot pools, then bring every block to CPU (undo any
|
||
pre-warm done by the streaming loader) so the engine's slot pools
|
||
are the sole owner of block placement. Finally place ``prewarmed``
|
||
blocks onto GPU slots (resident mode); the rest stay in home slots.
|
||
Lazy-mode blocks are left for on-demand disk reads when first used.
|
||
"""
|
||
self._xfer.cancel_all()
|
||
self._xfer.reset()
|
||
|
||
# Undo any pipeline pre-warm: every block starts on CPU.
|
||
for block in self._blocks:
|
||
try:
|
||
if any(p.is_cuda for p in block.parameters()):
|
||
move_module_to(block, "cpu")
|
||
except Exception:
|
||
pass
|
||
|
||
prewarmed = self._window.prewarmed
|
||
for i, block in enumerate(self._blocks):
|
||
if i < prewarmed:
|
||
self._xfer.init_gpu(i, block)
|
||
elif isinstance(self._home, RamHome):
|
||
self._xfer.init_home(i, block)
|
||
# DiskHome off-GPU blocks: left for lazy load on first use.
|
||
|
||
# Ensure peripherals start on CPU (model's home device).
|
||
for mod in self._peri.values():
|
||
if mod is not None:
|
||
move_module_to(mod, "cpu")
|
||
self._budget.maybe_flush(self._xfer.device, reserve_blocks=0)
|
||
|
||
# Pin entries are NOT pre-built here — they are built lazily inside
|
||
# ``load_block`` AFTER ``ensure_home`` has materialised the block's
|
||
# weights on CPU, guaranteeing the ``SlotEntry`` keys match the real
|
||
# parameter set (fp8 deferred-weight layers are only registered after
|
||
# ``load_state_dict`` runs inside ``ensure_ram``).
|
||
|
||
# ── Fold LoRA on GPU ring ─────────────────────────────────────
|
||
# All folding happens on the GPU ring (in-place), never on the CPU
|
||
# home pool. Prewarmed blocks fold here; blocks loaded later fold in
|
||
# prepare() right after load_block. "Folded or not" is inferred from
|
||
# whether each slot still carries its co-located LoRA (see
|
||
# _slot_has_lora) — no external _lora_folded set is needed, so a slot
|
||
# that was offloaded (and wrote its merged weights home) is simply
|
||
# skipped on re-load.
|
||
_folded = 0
|
||
for bidx in list(self._xfer._block_gpu):
|
||
if self._slot_has_lora(bidx):
|
||
self._fold_lora_on_gpu(bidx)
|
||
_folded += 1
|
||
if _folded:
|
||
logger.info("Folded LoRA on GPU ring: %d / %d blocks",
|
||
_folded, len(self._blocks))
|
||
|
||
# ------------------------------------------------------------------
|
||
# Public API — called from the model's forward pass
|
||
# ------------------------------------------------------------------
|
||
|
||
def prepare(self, block_idx: int) -> None:
|
||
"""Ensure block *block_idx* and its window are on GPU.
|
||
|
||
Offloads blocks that have fallen outside the window and loads newly
|
||
needed blocks. Off-GPU blocks live in the fixed CPU home pool; the
|
||
swap engine copies them to/from GPU slots with ``copy_`` — no new
|
||
tensors are allocated, so host RAM stays flat at ``(N-W)``. In lazy
|
||
mode the home pool is populated from disk on demand (inside the
|
||
engine's ``ensure_home``).
|
||
|
||
Call at the top of each loop iteration *before* using the block.
|
||
"""
|
||
if block_idx == 0:
|
||
self.prepare_blocks_entry()
|
||
|
||
to_offload = sorted(self._window.to_offload(block_idx))
|
||
to_load_now = sorted(self._window.to_load(block_idx))
|
||
did_offload = bool(to_offload)
|
||
|
||
# Interleave load (home→GPU) and offload (GPU→home). Crucially we
|
||
# *load first*, so the incoming block frees its home slot before the
|
||
# outgoing block needs one — the slot pool stays balanced at (N-W)
|
||
# and host RAM never spikes toward the whole model. Both sides use
|
||
# the recycled slot pools, so the transient VRAM/RAM churn is bounded
|
||
# to ~one block.
|
||
#
|
||
# A trailing block whose offload finds no free home slot is *deferred*
|
||
# (left on GPU) rather than raising — at the tail of a pass the window
|
||
# has shrunk below W, so more than N-W blocks would be off-GPU; the
|
||
# deferred blocks simply linger on GPU and are fully evicted at
|
||
# ``evict_all``. This is what keeps RAM flat at (N-W) throughout.
|
||
n = max(len(to_offload), len(to_load_now))
|
||
for k in range(n):
|
||
if k < len(to_load_now):
|
||
j = to_load_now[k]
|
||
if not self._xfer.sync_prefetch(j):
|
||
# Async load on the transfer stream, then sync the
|
||
# just-recorded event so the block's weights are fully
|
||
# resident on GPU before forward. Using the transfer
|
||
# stream (``non_blocking=True``) lets the copy overlap
|
||
# with any remaining offload work.
|
||
self._xfer.load_block(j, self._blocks[j], non_blocking=True)
|
||
self._xfer.sync_prefetch(j) # wait for async H2D
|
||
self._window.mark_loaded(j)
|
||
# Fold on GPU ring right after the block lands there (works for
|
||
# both RamHome and lazy DiskHome; prefetched blocks land on GPU
|
||
# via start_prefetch and are folded here too). _slot_has_lora
|
||
# skips a block whose slot already holds merged weights (from a
|
||
# previous offload), so re-folding never happens.
|
||
if self._slot_has_lora(j):
|
||
self._fold_lora_on_gpu(j)
|
||
if k < len(to_offload):
|
||
i = to_offload[k]
|
||
if self._xfer.offload_block(i, self._blocks[i]):
|
||
self._window.mark_offloaded(i)
|
||
# deferred -> block stays on GPU; retry on the next prepare
|
||
|
||
# Bulletproof residency guarantee.
|
||
for j in sorted(self._window.needed(block_idx)):
|
||
if self._xfer._block_gpu.get(j) is None:
|
||
if not self._xfer.sync_prefetch(j):
|
||
self._xfer.load_block(j, self._blocks[j], non_blocking=True)
|
||
self._xfer.sync_prefetch(j)
|
||
self._window.mark_loaded(j)
|
||
if self._slot_has_lora(j):
|
||
self._fold_lora_on_gpu(j)
|
||
|
||
if did_offload:
|
||
self._budget.maybe_flush(self._xfer.device, reserve_blocks=2)
|
||
|
||
def prefetch_next(self, block_idx: int) -> None:
|
||
"""Start async H2D copies for blocks just beyond the window.
|
||
|
||
Call after ``prepare(block_idx)`` and before the compute for
|
||
*block_idx*. Copies run on the transfer stream in parallel with
|
||
the current block's forward pass. In lazy mode the home pool is
|
||
populated from disk as a side effect of the load. No-op when
|
||
prefetch is disabled.
|
||
"""
|
||
if not self._prefetch:
|
||
return
|
||
start = block_idx + self._window.window_size
|
||
if start >= self._window.total:
|
||
return
|
||
count = self._max_prefetch(start)
|
||
end = min(self._window.total, start + count)
|
||
for i in range(start, end):
|
||
if self._window.is_on_gpu(i):
|
||
continue
|
||
self._xfer.start_prefetch(i, self._blocks[i])
|
||
|
||
def prepare_pre_forward(self) -> None:
|
||
"""Ensure patch_embedding is on GPU before ``pre_forward``."""
|
||
self._xfer.cancel_all()
|
||
self._load_peripheral("patch_embedding")
|
||
|
||
def prepare_blocks_entry(self) -> None:
|
||
"""Ensure embeddings/projection are on GPU for ``transformer_forward``."""
|
||
for name in ("text_embedding", "time_embedding", "time_projection"):
|
||
self._load_peripheral(name)
|
||
|
||
def prepare_head(self) -> None:
|
||
"""Ensure head is on GPU before final unpatchify."""
|
||
self._load_peripheral("head")
|
||
|
||
def is_on_gpu(self, block_idx: int) -> bool:
|
||
"""Check if a block is currently on GPU (for S² drop / STG skip)."""
|
||
return self._window.is_on_gpu(block_idx)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Lifecycle
|
||
# ------------------------------------------------------------------
|
||
|
||
def evict_all(self) -> None:
|
||
"""Move all blocks and peripherals back to CPU home slots.
|
||
|
||
For lazy mode, also free the home-pool CPU RAM (disk is the source
|
||
of truth, so the buffers are not needed until the next load).
|
||
"""
|
||
self._xfer.cancel_all()
|
||
self._xfer.sync_all()
|
||
lazy = self._home is not None and not isinstance(self._home, RamHome)
|
||
for i in sorted(self._window.on_gpu):
|
||
if lazy:
|
||
# Disk is the source of truth and every home slot is about to
|
||
# be released + the pools cleared below, so writing GPU slots
|
||
# back would be pure waste (~window × block-size of D2H on
|
||
# every expert switch / unload). Drop the mapping only.
|
||
self._xfer.discard_block(i)
|
||
else:
|
||
self._xfer.offload_block(i, self._blocks[i], force=True)
|
||
self._window.clear()
|
||
for name in self._PERIPHERAL_NAMES:
|
||
self._offload_peripheral(name)
|
||
|
||
# Drain every async D2H issued above BEFORE releasing RAM or
|
||
# clearing pools: the DMA must not still be reading GPU slots or
|
||
# writing home slots once their references are dropped (freed host
|
||
# pages could be reused by an unrelated tensor mid-copy).
|
||
self._xfer.sync_all()
|
||
self._xfer._d2h_events.clear()
|
||
|
||
# Lazy mode: drop the in-RAM weights (disk-backed) and let the engine
|
||
# re-populate the home slot on the next load. Resident mode keeps
|
||
# them — they are the model's persistent CPU copy.
|
||
if lazy:
|
||
for i in range(self._window.total):
|
||
self._home.release(i)
|
||
self._xfer.forget_home(i)
|
||
self._home.join()
|
||
|
||
try:
|
||
torch.cuda.empty_cache()
|
||
except Exception:
|
||
pass
|
||
# Drop all pool tensor references — without this the home & ring
|
||
# buffer SlotEntry tensors keep RAM/VRAM pinned across model switches.
|
||
self._xfer.clear_pools()
|
||
|
||
def shutdown(self) -> None:
|
||
"""Deterministic full teardown, safe to call on model unload/switch.
|
||
|
||
Order matters on Windows: stop the disk workers *first* and let them
|
||
finish (so no thread is still writing block weights into the model),
|
||
then synchronise every CUDA stream, then evict everything.
|
||
"""
|
||
try:
|
||
self._xfer.cancel_all()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
self._home.join()
|
||
except Exception:
|
||
pass
|
||
self._xfer.sync_all()
|
||
self.evict_all()
|
||
try:
|
||
self._home.shutdown()
|
||
except Exception:
|
||
pass
|
||
# Drop all pool tensor references and block list so the old model's
|
||
# Parameter objects have zero live references -> GC reclaims RAM+VRAM.
|
||
self._xfer.clear_pools()
|
||
self._blocks = []
|
||
self._home = None
|
||
|
||
# ------------------------------------------------------------------
|
||
# Internal helpers
|
||
# ------------------------------------------------------------------
|
||
|
||
def _load_peripheral(self, name: str) -> None:
|
||
mod = self._peri.get(name)
|
||
if mod is not None:
|
||
move_module_to(mod, self._xfer.device)
|
||
|
||
def _offload_peripheral(self, name: str) -> None:
|
||
mod = self._peri.get(name)
|
||
if mod is not None:
|
||
move_module_to(mod, "cpu")
|
||
|
||
def _max_prefetch(self, start: int) -> int:
|
||
"""How many blocks we can safely prefetch given free VRAM."""
|
||
if not self._xfer.prefetch or self._xfer._stream is None:
|
||
return 0
|
||
if self._budget.block_mb <= 0:
|
||
# A block group with no parameters (e.g. a norm-only slice)
|
||
# would divide by zero below; nothing to prefetch.
|
||
return 0
|
||
free_mb = self._budget.free_mb(self._xfer.device)
|
||
reserve_mb = self._budget.block_mb * 2
|
||
avail_mb = max(0.0, free_mb - reserve_mb)
|
||
max_by_mem = int(avail_mb / (self._budget.block_mb * 1.2))
|
||
return max(0, min(self._xfer.prefetch_count, max_by_mem))
|
||
|
||
# ── GPU LoRA folding (fused into ring-buffer load) ──────────────
|
||
|
||
def _fold_lora_on_entries(self, entries: dict, block_idx: int) -> None:
|
||
"""Fold any LoRA co-located in *entries* into the block weights in-place.
|
||
|
||
Reads ``slot.lora`` (set when the slot was materialised from the
|
||
block's ``_lora_payload``) instead of any global map. After folding,
|
||
``slot.lora`` is set to ``None`` — the slot becomes a pure ``(block)``
|
||
and re-loads skip folding automatically. "Folded or not" is therefore
|
||
inferred from the slot contents; no external bookkeeping set is needed.
|
||
"""
|
||
from ..utils.lora import _lora_delta, _requantize_fp8
|
||
from torch import float32 as _f32
|
||
|
||
for param_name, slot in entries.items():
|
||
lora_entries = slot.lora
|
||
if not lora_entries:
|
||
continue
|
||
|
||
# Work on the SAME device as the slot (GPU ring or CPU home),
|
||
# so LoRA tensors are moved there too — never to a mismatched device.
|
||
work_device = slot.data.device
|
||
|
||
# ── nvfp4: dual-scale dequant → fp32 fold → dual-scale requant ──
|
||
if slot.is_qt and "NVFP4" in slot.layout_cls:
|
||
from ..utils.lora import _apply_streaming_loras_nvfp4
|
||
new_q, new_scale, new_bs = _apply_streaming_loras_nvfp4(
|
||
slot.data, slot.scale, slot.extra["block_scale"],
|
||
slot.orig_dtype, slot.orig_shape, slot.meta, lora_entries)
|
||
slot.data.copy_(new_q)
|
||
slot.scale.copy_(new_scale)
|
||
slot.extra["block_scale"].copy_(new_bs)
|
||
slot.lora = None
|
||
continue
|
||
|
||
if slot.is_qt:
|
||
base_f = slot.data.to(_f32) * slot.scale.to(_f32)
|
||
is_fp8 = True
|
||
else:
|
||
base_f = slot.data.to(_f32)
|
||
is_fp8 = False
|
||
|
||
has_lora = any(e.get("A") is not None and e.get("B") is not None
|
||
for e in lora_entries)
|
||
has_dora = any(e.get("diff_b") is not None for e in lora_entries) or \
|
||
any(e.get("diff") is not None for e in lora_entries)
|
||
|
||
if not has_lora and not has_dora:
|
||
slot.lora = None
|
||
continue
|
||
|
||
if has_dora and base_f.dim() == 2:
|
||
init_norm = base_f.norm(dim=1, keepdim=True).clamp(min=1e-8)
|
||
else:
|
||
init_norm = None
|
||
|
||
if has_lora:
|
||
for entry in lora_entries:
|
||
A, B = entry.get("A"), entry.get("B")
|
||
if A is None or B is None:
|
||
continue
|
||
A_gpu = A.to(device=work_device, dtype=_f32, non_blocking=False)
|
||
B_gpu = B.to(device=work_device, dtype=_f32, non_blocking=False)
|
||
delta, _ = _lora_delta(
|
||
A_gpu, B_gpu,
|
||
entry.get("alpha"), entry.get("strength", 1.0),
|
||
base_f.shape,
|
||
)
|
||
base_f = base_f + delta
|
||
|
||
diff_b = None
|
||
diff = None
|
||
for entry in lora_entries:
|
||
if entry.get("diff_b") is not None:
|
||
diff_b = entry["diff_b"]
|
||
if entry.get("diff") is not None:
|
||
diff = entry["diff"]
|
||
|
||
if diff_b is not None and base_f.dim() == 2:
|
||
temp_norm = base_f.norm(dim=1, keepdim=True).clamp(min=1e-8)
|
||
diff_gpu = diff_b.to(device=work_device, dtype=_f32).reshape(-1, 1)
|
||
m = (init_norm + diff_gpu).clamp(min=0.0)
|
||
base_f = m * base_f / temp_norm
|
||
elif diff is not None and base_f.dim() == 1:
|
||
base_f = base_f + diff.to(device=work_device, dtype=_f32)
|
||
|
||
if is_fp8:
|
||
new_data, new_scale = _requantize_fp8(base_f)
|
||
slot.data.copy_(new_data)
|
||
slot.scale.copy_(new_scale)
|
||
else:
|
||
slot.data.copy_(base_f.to(slot.data.dtype))
|
||
|
||
# Consume: slot is now pure (block); re-loads skip folding.
|
||
slot.lora = None
|
||
|
||
def _fold_lora_on_gpu(self, block_idx: int) -> None:
|
||
"""Fold LoRA for *block_idx* on its GPU ring slot in-place (if present)."""
|
||
gslot = self._xfer._block_gpu.get(block_idx)
|
||
if gslot is None:
|
||
return
|
||
self._fold_lora_on_entries(self._xfer._gpu_pool[gslot], block_idx)
|
||
|
||
def _slot_has_lora(self, block_idx: int) -> bool:
|
||
"""Whether *block_idx*'s GPU ring slot still carries unfused LoRA.
|
||
|
||
Used to decide whether to fold on load — replaces the old
|
||
``_lora_folded`` bookkeeping set: a slot that has no LoRA is already
|
||
merged, so re-loads are skipped automatically.
|
||
"""
|
||
gslot = self._xfer._block_gpu.get(block_idx)
|
||
if gslot is None:
|
||
return False
|
||
return any(e.lora for e in self._xfer._gpu_pool[gslot].values())
|
||
|