trainstream 0.2.0__tar.gz → 0.3.0__tar.gz

{trainstream-0.2.0/trainstream.egg-info → trainstream-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trainstream
-Version: 0.2.0
+Version: 0.3.0
 Summary: Zero-dependency PCIe offloading engine for PyTorch training (optimizer, gradients, activations).
 Home-page: https://github.com/johnanishg/TrainStream
 Author: John Anish G
@@ -58,7 +58,7 @@ TrainStream is a C++/CUDA extension for PyTorch. It keeps only the weights in VR
 | 1.8B GPT, weights paged to RAM (0.2) | — | 2.4 GiB peak, 2.1 s/step; 1.84 GiB inside a 4 GB VRAM cap with bf16 moments |
 | 3.6B GPT full fine-tune (all parameters trained) | needs ~58 GB of VRAM | 3.1 GiB peak, 3.15 s/step, 30 GiB of host RAM |
 | LoRA r=16 on the Qwen2.5-7B architecture (random weights) | 20.3 GB peak, 0.27 s/step | 4.6 GB peak, 1.25 s/step |
-| LoRA r=16 on the Qwen2.5-14B architecture (random weights; 29.5 GB) | does not fit | 5.6 GB peak, 2.48 s/step, 36 GB of host RAM |
+| LoRA r=16 on Qwen2.5-14B (29.5 GB of weights) | does not fit | 1.37 s/step: the memory planner keeps 28 of 48 layers in VRAM (19.9 GiB peak) and streams 20 from RAM |
 | 85M GPT (fits either way) | 2.55 GiB, 43.7 ms/step | `offload="auto"` switches to stock fused AdamW: 2.56 GiB, 43.8 ms/step |
 
 Offloading costs throughput when a model already fits, which is why `offload="auto"` measures two steps and then switches to plain fused GPU AdamW if offloading isn't needed. TrainStream's value is training what does not fit, or freeing VRAM for larger batches. All numbers come from a single machine (RTX 5090 Laptop, Windows 11).
@@ -67,7 +67,7 @@ Offloading costs throughput when a model already fits, which is why `offload="au
 
 - An NVIDIA GPU and a CUDA build of PyTorch
 - The CUDA toolkit matching PyTorch's CUDA major version, and a C++ compiler (Visual Studio 2022+ Build Tools with "Desktop development with C++" on Windows, GCC 11+ on Linux)
-- Tested on Windows 11 (PyTorch 2.14.1, CUDA 13.4, MSVC, Python 3.14) and Ubuntu 24.04 under WSL2 (PyTorch 2.14.1+cu130, CUDA 13.3, GCC 13.3, Python 3.12): all 78 tests pass on both, RTX 5090 Laptop GPU. Native Linux and other GPUs are untested.
+- Tested on Windows 11 (PyTorch 2.14.1, CUDA 13.4, MSVC, Python 3.14) and Ubuntu 24.04 under WSL2 (PyTorch 2.14.1+cu130, CUDA 13.3, GCC 13.3, Python 3.12): all 89 tests pass on both, RTX 5090 Laptop GPU. Native Linux and other GPUs are untested.
 
 ## Install
 
@@ -118,9 +118,10 @@ model, optimizer = trainstream.wrap(model, optimizer, device="cuda",
 | `wrap()` option | What it does |
 |---|---|
 | `offload="auto"` | switches to stock fused AdamW after 2 steps if the model fits (`True` always offloads) |
-| `param_offload="auto"` | pages block weights to host RAM when they'd take more than half the GPU; `True`, a fraction, or block indices |
+| `param_offload="auto"` | memory planner: measures the VRAM and RAM of the machine it runs on, keeps as many transformer blocks resident as fit, pages the rest, and refines the split from the first steps' measured peaks; `True`, a fraction, or block indices fix it |
 | `activation_offload="auto"` | pages a profiled share of activations (`True`, a fraction, `False`) |
-| `moment_dtype="bfloat16"` | bf16 optimizer moments with stochastic rounding: half the host memory for them |
+| `moment_dtype="auto"` | fp32 optimizer moments, or bf16 (stochastic rounding, half the memory) when fp32 wouldn't fit in RAM |
+| `check_memory=True` | refuse at `wrap()` with a clear explanation when the run can't fit in RAM/VRAM |
 | `distributed` | data parallel when a `torch.distributed` group with more than 1 rank exists; gradients averaged on the host over gloo |
 | `max_grad_norm`, `accumulation_steps` | clipping and gradient accumulation, done inside the optimizer |
 

{trainstream-0.2.0 → trainstream-0.3.0}/README.md

@@ -18,7 +18,7 @@ TrainStream is a C++/CUDA extension for PyTorch. It keeps only the weights in VR
 | 1.8B GPT, weights paged to RAM (0.2) | — | 2.4 GiB peak, 2.1 s/step; 1.84 GiB inside a 4 GB VRAM cap with bf16 moments |
 | 3.6B GPT full fine-tune (all parameters trained) | needs ~58 GB of VRAM | 3.1 GiB peak, 3.15 s/step, 30 GiB of host RAM |
 | LoRA r=16 on the Qwen2.5-7B architecture (random weights) | 20.3 GB peak, 0.27 s/step | 4.6 GB peak, 1.25 s/step |
-| LoRA r=16 on the Qwen2.5-14B architecture (random weights; 29.5 GB) | does not fit | 5.6 GB peak, 2.48 s/step, 36 GB of host RAM |
+| LoRA r=16 on Qwen2.5-14B (29.5 GB of weights) | does not fit | 1.37 s/step: the memory planner keeps 28 of 48 layers in VRAM (19.9 GiB peak) and streams 20 from RAM |
 | 85M GPT (fits either way) | 2.55 GiB, 43.7 ms/step | `offload="auto"` switches to stock fused AdamW: 2.56 GiB, 43.8 ms/step |
 
 Offloading costs throughput when a model already fits, which is why `offload="auto"` measures two steps and then switches to plain fused GPU AdamW if offloading isn't needed. TrainStream's value is training what does not fit, or freeing VRAM for larger batches. All numbers come from a single machine (RTX 5090 Laptop, Windows 11).
@@ -27,7 +27,7 @@ Offloading costs throughput when a model already fits, which is why `offload="au
 
 - An NVIDIA GPU and a CUDA build of PyTorch
 - The CUDA toolkit matching PyTorch's CUDA major version, and a C++ compiler (Visual Studio 2022+ Build Tools with "Desktop development with C++" on Windows, GCC 11+ on Linux)
-- Tested on Windows 11 (PyTorch 2.14.1, CUDA 13.4, MSVC, Python 3.14) and Ubuntu 24.04 under WSL2 (PyTorch 2.14.1+cu130, CUDA 13.3, GCC 13.3, Python 3.12): all 78 tests pass on both, RTX 5090 Laptop GPU. Native Linux and other GPUs are untested.
+- Tested on Windows 11 (PyTorch 2.14.1, CUDA 13.4, MSVC, Python 3.14) and Ubuntu 24.04 under WSL2 (PyTorch 2.14.1+cu130, CUDA 13.3, GCC 13.3, Python 3.12): all 89 tests pass on both, RTX 5090 Laptop GPU. Native Linux and other GPUs are untested.
 
 ## Install
 
@@ -78,9 +78,10 @@ model, optimizer = trainstream.wrap(model, optimizer, device="cuda",
 | `wrap()` option | What it does |
 |---|---|
 | `offload="auto"` | switches to stock fused AdamW after 2 steps if the model fits (`True` always offloads) |
-| `param_offload="auto"` | pages block weights to host RAM when they'd take more than half the GPU; `True`, a fraction, or block indices |
+| `param_offload="auto"` | memory planner: measures the VRAM and RAM of the machine it runs on, keeps as many transformer blocks resident as fit, pages the rest, and refines the split from the first steps' measured peaks; `True`, a fraction, or block indices fix it |
 | `activation_offload="auto"` | pages a profiled share of activations (`True`, a fraction, `False`) |
-| `moment_dtype="bfloat16"` | bf16 optimizer moments with stochastic rounding: half the host memory for them |
+| `moment_dtype="auto"` | fp32 optimizer moments, or bf16 (stochastic rounding, half the memory) when fp32 wouldn't fit in RAM |
+| `check_memory=True` | refuse at `wrap()` with a clear explanation when the run can't fit in RAM/VRAM |
 | `distributed` | data parallel when a `torch.distributed` group with more than 1 rank exists; gradients averaged on the host over gloo |
 | `max_grad_norm`, `accumulation_steps` | clipping and gradient accumulation, done inside the optimizer |
 

{trainstream-0.2.0 → trainstream-0.3.0}/setup.py

@@ -66,7 +66,7 @@ with open(os.path.join(HERE, "README.md"), encoding="utf-8") as fh:
 
 setup(
     name="trainstream",
-    version="0.2.0",
+    version="0.3.0",
     author="John Anish G",
     license="Apache-2.0",
     description="Zero-dependency PCIe offloading engine for PyTorch training (optimizer, gradients, activations).",

trainstream-0.3.0/tests/test_planner.py

@@ -0,0 +1,247 @@
+"""The memory planner: decisions from the measured VRAM / RAM, and live re-planning."""
+
+import copy
+import json
+import os
+import subprocess
+import sys
+import textwrap
+
+import pytest
+import torch
+from torch import nn
+
+import trainstream
+from trainstream.planner import Budget, MemoryPlanner, spread
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+GiB = 1 << 30
+DEV = "cuda"
+HP = dict(lr=2e-3, betas=(0.9, 0.95), weight_decay=0.05)
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.dirname(HERE)
+
+
+def budget(vram_gib, ram_gib=64):
+    return Budget(vram_total=int(vram_gib * GiB), vram_usable=int(vram_gib * GiB), ram_total=int(ram_gib * GiB),
+                  ram_available=int(ram_gib * GiB), pin_budget=int(0.4 * ram_gib * GiB))
+
+
+# --------------------------------------------------------------------------- decisions
+def test_spread_is_even_and_exact():
+    assert spread(0, list(range(10))) == []
+    assert spread(10, list(range(10))) == list(range(10))
+    s = spread(4, list(range(16)))
+    assert len(s) == 4 and all(b - a == 4 for a, b in zip(s, s[1:]))
+    assert len(spread(7, list(range(48)))) == 7
+
+
+@pytest.mark.parametrize("vram,expect_paged", [(80, 0), (24, None), (8, None), (4, None)])
+def test_static_plan_follows_the_vram_budget(vram, expect_paged):
+    """48 blocks of 0.55 GiB (Qwen2.5-14B in bf16) and 3 GiB outside them."""
+    blocks = [int(0.55 * GiB)] * 48
+    p = MemoryPlanner(budget(vram), blocks, fixed_vram=3 * GiB, inflight_bytes=GiB // 2)
+    paged = p.initial_paged()
+    if expect_paged == 0:
+        assert paged == set()  # an 80 GB card holds the whole model
+        return
+    resident = 48 - len(paged)
+    usable = 0.9 * vram * GiB
+    assert 3 * GiB + resident * 0.55 * GiB < usable  # what stays resident fits
+    if vram >= 24:
+        assert resident >= 20  # most of a 24 GB card is used for weights
+    assert len(paged) == len(set(paged)) and all(0 <= i < 48 for i in paged)
+
+
+def test_refinement_moves_blocks_in_with_headroom_and_out_when_full():
+    blocks = [GiB // 2] * 20
+    p = MemoryPlanner(budget(16), blocks, fixed_vram=2 * GiB, inflight_bytes=GiB // 4)
+    paged = set(range(0, 20, 2))  # 10 paged
+    grown = p.refine(peak=6 * GiB, paged=paged)  # lots of headroom
+    assert grown is not None and grown < paged
+    shrunk = p.refine(peak=int(15.5 * GiB), paged=grown)  # too close to the budget
+    assert shrunk is not None and shrunk > grown
+    p.refine(peak=int(13.5 * GiB), paged=shrunk)
+    p.refine(peak=int(13.5 * GiB), paged=shrunk)
+    assert p.done  # two steady steps settle the plan
+
+
+def test_wrap_refuses_what_cannot_fit_in_ram(monkeypatch):
+    import trainstream.planner as planner
+    monkeypatch.setattr(planner, "available_ram_bytes", lambda: 64 << 20)  # 64 MiB "free"
+    model = nn.Sequential(*[nn.Linear(1024, 1024) for _ in range(8)]).to(DEV)
+    with pytest.raises(MemoryError, match="host RAM"):
+        trainstream.wrap(model, torch.optim.AdamW(model.parameters()), offload=True, activation_offload=False)
+    # the explicit opt-out still wraps
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters()), offload=True,
+                                  activation_offload=False, check_memory=False)
+    trainstream.unwrap(model)
+
+
+def test_moment_precision_follows_available_ram(monkeypatch):
+    import trainstream.planner as planner
+    model = nn.Sequential(*[nn.Linear(1024, 1024) for _ in range(8)]).to(DEV)
+    numel = sum(p.numel() for p in model.parameters())
+    # host need: fp32 master 4 + staging 4 + moments 2x(4 or 2) B/param -> 16 (fp32) vs 12 (bf16); give 14
+    room = int(numel * 14 / 0.9)
+    monkeypatch.setattr(planner, "available_ram_bytes", lambda: room)
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters()), offload=True, activation_offload=False)
+    try:
+        assert opt.moment_dtype == "bfloat16"
+    finally:
+        trainstream.unwrap(model)
+    monkeypatch.setattr(planner, "available_ram_bytes", lambda: 64 << 30)
+    model = nn.Sequential(*[nn.Linear(1024, 1024) for _ in range(8)]).to(DEV)
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters()), offload=True, activation_offload=False)
+    try:
+        assert opt.moment_dtype == "float32"
+    finally:
+        trainstream.unwrap(model)
+
+
+# --------------------------------------------------------------------------- live re-planning
+class Block(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.ln = nn.LayerNorm(d)
+        self.fc1, self.fc2 = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)
+
+    def forward(self, x):
+        return x + self.fc2(torch.nn.functional.gelu(self.fc1(self.ln(x))))
+
+
+class Stack(nn.Module):
+    def __init__(self, d=256, n=6, vocab=128):
+        super().__init__()
+        self.emb = nn.Embedding(vocab, d)
+        self.blocks = nn.ModuleList(Block(d) for _ in range(n))
+        self.head = nn.Linear(d, vocab)
+
+    def forward(self, idx):
+        x = self.emb(idx)
+        for b in self.blocks:
+            x = b(x)
+        return self.head(x)
+
+
+def lm_steps(model, opt, data, between=None):
+    out = []
+    for k, t in enumerate(data):
+        logits = model(t[:, :-1])
+        loss = nn.functional.cross_entropy(logits.float().reshape(-1, logits.size(-1)), t[:, 1:].reshape(-1))
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+        out.append(loss.item())
+        if between:
+            between(k)
+    return out
+
+
+def tokens(n, seed=3, batch=8, seq=32, vocab=128):
+    g = torch.Generator().manual_seed(seed)
+    return [torch.randint(0, vocab, (batch, seq + 1), generator=g).to(DEV) for _ in range(n)]
+
+
+def test_blocks_move_between_vram_and_ram_mid_training_without_changing_results():
+    torch.manual_seed(0)
+    ref = Stack().to(DEV)
+    mine = copy.deepcopy(ref)
+    opt_ref = torch.optim.AdamW(ref.parameters(), foreach=False, **HP)
+    mine, opt = trainstream.wrap(mine, torch.optim.AdamW(mine.parameters(), **HP), offload=True,
+                                 param_offload=[0, 2, 4], activation_offload=False)
+    pager = trainstream.wrapper.unwrap_session(mine).pager
+    moves = {1: ([0], [1, 5]), 3: ([1, 2, 4, 5], []), 5: ([], [0, 3])}  # after step k: (unpage, page)
+
+    def between(k):
+        if k in moves:
+            unpage, page = moves[k]
+            for i in page:
+                pager.page(i)
+            for i in unpage:
+                pager.unpage(i)
+
+    try:
+        data = tokens(8)
+        a = lm_steps(ref, opt_ref, data)
+        b = lm_steps(mine, opt, data, between)
+        opt.synchronize()
+        assert b == pytest.approx(a, abs=1e-4)
+        assert pager.paged == {0, 3}
+        sizes = [sum(p.untyped_storage().nbytes() for p in blk.parameters()) for blk in mine.blocks]
+        assert all(sizes[i] > 0 for i in (1, 2, 4, 5)) and sizes[3] == 0  # resident blocks hold weights
+        sd, sd_ref = mine.state_dict(), ref.state_dict()
+        assert max((sd[k].float().cpu() - sd_ref[k].float().cpu()).abs().max().item() for k in sd_ref) < 2e-4
+    finally:
+        trainstream.unwrap(mine)
+
+
+CAP_SCRIPT = textwrap.dedent(r"""
+    import json, sys, torch
+    sys.path[:0] = [ROOT, HERE]
+    import trainstream
+    from test_planner import Stack, tokens, lm_steps, HP
+    cap = CAP
+    torch.cuda.set_per_process_memory_fraction(cap / torch.cuda.get_device_properties(0).total_memory, 0)
+    torch.manual_seed(0)
+    model = Stack(d=1024, n=8)  # 8 blocks of 50 MB
+    opt = torch.optim.AdamW(model.parameters(), **HP)
+    model, opt = trainstream.wrap(model, opt, device="cuda", offload=True)
+    s = trainstream.wrapper.unwrap_session(model)
+    first = sorted(s.pager.paged)
+    peaks = []
+    def between(k):
+        peaks.append(torch.cuda.max_memory_allocated())
+    losses = lm_steps(model, opt, tokens(14, batch=4, seq=64), between)
+    opt.synchronize()
+    print("@@" + json.dumps(dict(first=first, final=sorted(s.pager.paged), settled=s.planner.done, losses=losses,
+                                 events=s.events, moments=opt.moment_dtype, cap=cap)))
+""")
+
+
+def test_auto_plan_under_a_vram_cap_finds_the_split_and_trains_correctly():
+    cap = 450 << 20  # 8 blocks of 50 MB do not all fit next to the rest
+    code = CAP_SCRIPT.replace("ROOT", repr(ROOT)).replace("HERE", repr(HERE)).replace("CAP", str(cap))
+    out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600)
+    assert out.returncode == 0, out.stdout + out.stderr
+    res = json.loads([l for l in out.stdout.splitlines() if l.startswith("@@")][-1][2:])
+    assert res["first"], "a 450 MB cap must start with some blocks paged"
+    assert res["settled"]
+    assert res["final"] != res["first"]  # measured peaks changed the plan
+    assert 0 < len(res["final"]) < len(res["first"])  # measured headroom moved some blocks into VRAM
+    assert any("memory plan settled" in e for e in res["events"])
+    # Moving blocks between VRAM and RAM must not change the arithmetic at all: the
+    # capped, re-planned run equals an uncapped, unpaged TrainStream run exactly.
+    # (Against torch.optim.AdamW this lr is unstable enough to amplify last-bit
+    # differences after ~8 steps, so that comparison lives in the smaller tests.)
+    torch.manual_seed(0)
+    ref = Stack(d=1024, n=8)
+    ref, opt = trainstream.wrap(ref, torch.optim.AdamW(ref.parameters(), **HP), device=DEV, offload=True,
+                                param_offload=False)
+    try:
+        a = lm_steps(ref, opt, tokens(14, batch=4, seq=64))
+    finally:
+        trainstream.unwrap(ref)
+    assert res["losses"] == a
+
+
+def test_watchdog_pages_a_block_when_a_later_step_grows():
+    torch.manual_seed(0)
+    model = Stack().to(DEV)
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters(), **HP), offload=True,
+                                  param_offload=[0], activation_offload=False)
+    s = trainstream.wrapper.unwrap_session(model)
+    from trainstream.planner import Budget, MemoryPlanner
+    used = torch.cuda.memory_allocated()
+    s.planner = MemoryPlanner(Budget(vram_total=used + (64 << 20), vram_usable=used + (64 << 20), ram_total=1 << 36,
+                                     ram_available=1 << 35, pin_budget=1 << 34),
+                              s.pager.block_bytes, fixed_vram=0, inflight_bytes=1 << 20)
+    s.planner.done = True  # settled; now a step with a much larger batch arrives
+    try:
+        before = len(s.pager.paged)
+        lm_steps(model, opt, tokens(1, batch=64, seq=128))
+        assert len(s.pager.paged) == before + 1
+        assert any("memory watchdog" in e for e in s.events)
+    finally:
+        trainstream.unwrap(model)

{trainstream-0.2.0 → trainstream-0.3.0}/trainstream/__init__.py

@@ -33,7 +33,7 @@ from .paging import ParamPager
 from .profiler import HardwareProfile, profile_hardware
 from .wrapper import TrainStreamSession, memory_report, unwrap, wrap
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 __all__ = [
     "wrap",

{trainstream-0.2.0 → trainstream-0.3.0}/trainstream/paging.py

@@ -6,16 +6,25 @@ block the VRAM storage of its weights is released (``storage.resize_(0)``, the
 same mechanism FSDP uses), and the pager streams the host copy back in on the
 H2D stream whenever the block is about to run:
 
-  forward   block i: load i (usually prefetched), prefetch i+1; release i
-            after it ran (the last block stays, backward starts there)
-  backward  a hook on block i's output gradient loads i and prefetches i-1;
-            blocks above i are released — their input gradients, the last
-            computation that reads their weights, are already queued
-
-Frozen weights (LoRA base layers) get their own pinned host copy. Under
-torch.autocast the cast-weight cache is cleared after each paged block, or it
-would keep a bf16 copy of every paged weight alive until the region exits. Updates of
-paged parameters are never uploaded by the optimizer; the next load reads the
+  forward   block i: load i if it is paged (usually prefetched), prefetch the
+            next paged block(s); release i after it ran (the last paged block
+            stays, backward starts there)
+  backward  a hook on block i's output gradient loads i (if paged) and
+            prefetches the previous paged block(s); paged blocks above i are
+            released — their input gradients, the last computation that reads
+            their weights, are already queued
+
+Which blocks are paged can change between steps (``page`` / ``unpage``): the
+memory planner keeps as many blocks resident as VRAM allows. Hooks sit on every
+block, so a resident block prefetches the next paged one and transfers overlap
+with compute even when paged and resident blocks are interleaved.
+
+Frozen weights (LoRA base layers) of a paged block get a host copy of their
+own: pinned while that fits the page-lock budget, else ordinary RAM. It is
+dropped again when the block becomes resident. Under torch.autocast the
+cast-weight cache is cleared after each block, or it would keep a bf16 copy of
+every weight alive until the autocast region exits. Updates of paged
+parameters are never uploaded by the optimizer; the next load reads the
 updated host copy after the optimizer has finished with it, and the next
 update waits for that DMA to finish reading.
 """
@@ -69,65 +78,126 @@ def pageable_params(model: nn.Module, blocks: Sequence[nn.Module], indices: Iter
 
 
 class ParamPager:
-    def __init__(self, model: nn.Module, blocks: Sequence[nn.Module], optimizer, paged: Dict[int, List[nn.Parameter]],
-                 *, frozen_init: Optional[Dict[int, torch.Tensor]] = None, prefetch_depth: int = 1):
+    def __init__(self, model: nn.Module, blocks: Sequence[nn.Module], optimizer, paged: Iterable[int] = (), *,
+                 candidates: Optional[Dict[int, List[nn.Parameter]]] = None,
+                 frozen_init: Optional[Dict[int, torch.Tensor]] = None, prefetch_depth: int = 1):
         self.model = model
         self.blocks = list(blocks)
         self.n = len(self.blocks)
         self.opt = optimizer
         self.depth = max(1, int(prefetch_depth))
-        self.paged = {i: ps for i, ps in paged.items() if ps}
-        self.last = max(self.paged) if self.paged else -1
-        frozen_init = frozen_init if frozen_init is not None else {}
         self.device = optimizer.device
+        if candidates is None:
+            candidates = pageable_params(model, self.blocks, range(self.n))
         self._entries: Dict[int, list] = {}
+        for i in range(self.n):
+            self._entries[i] = [(p, optimizer._slot_of.get(id(p)), p.numel() * p.element_size())
+                                for p in candidates.get(i, [])]
+        self.block_bytes = [sum(nb for _, _, nb in self._entries[i]) for i in range(self.n)]
         self._host_frozen: Dict[int, torch.Tensor] = {}
         ram = physical_ram_bytes()
-        pin_budget = 0.40 * ram if ram else float("inf")
+        self.pin_budget = 0.40 * ram if ram else float("inf")
         self.frozen_pinned_bytes = 0
         self.frozen_pageable_bytes = 0
-        for i, ps in self.paged.items():
-            entries = []
-            for p in ps:
-                slot = optimizer._slot_of.get(id(p))
-                if slot is None:
-                    # Frozen weight: pinned host copy while that fits the page-lock
-                    # budget, else ordinary RAM (the CPU tensor it was built in is
-                    # reused as is). Popped one by one, so the model never exists
-                    # twice in host memory.
-                    src = frozen_init.pop(id(p), None)
-                    nbytes = p.numel() * p.element_size()
-                    if _C.pinned_stats()["live_bytes"] + nbytes <= pin_budget:
-                        host = _C.pinned_empty([p.numel()], p.dtype, "param.frozen", False, False)
-                        host.copy_(src.reshape(-1) if src is not None else p.detach().reshape(-1))
-                        self.frozen_pinned_bytes += nbytes
-                    else:
-                        host = (src.reshape(-1) if src is not None else p.detach().reshape(-1).cpu())
-                        self.frozen_pageable_bytes += nbytes
-                    del src
-                    self._host_frozen[id(p)] = host
-                entries.append((p, slot, p.numel() * p.element_size()))
-            self._entries[i] = entries
-        self.allocated = {i: True for i in self.paged}
-        self.fresh = {i: False for i in self.paged}
+        self.paged: Set[int] = set()
+        self._sorted: List[int] = []
+        self.allocated = {i: True for i in range(self.n)}
+        self.fresh = {i: False for i in range(self.n)}
         self.events: Dict[int, object] = {}
-        self.stats = dict(loads=0, load_bytes=0, releases=0)
-        self.bytes_paged = sum(nb for es in self._entries.values() for _, _, nb in es)
-        for i in self.paged:
-            self.release(i)
+        self.stats = dict(loads=0, load_bytes=0, releases=0, pages=0, unpages=0)
+        frozen_init = frozen_init if frozen_init is not None else {}
+        for i in sorted(set(paged)):
+            self.page(i, frozen_init)
         self._handles = [model.register_forward_pre_hook(self._root_pre)]
-        for i in self.paged:
-            b = self.blocks[i]
+        for i, b in enumerate(self.blocks):
             self._handles.append(b.register_forward_pre_hook(self._make_pre(i)))
             self._handles.append(b.register_forward_hook(self._make_post(i), always_call=True))
             self._handles.append(b._register_state_dict_hook(self._make_state_dict_hook(i)))
             self._handles.append(b._register_load_state_dict_pre_hook(self._make_load_pre(i), with_module=True))
             self._handles.append(b.register_load_state_dict_post_hook(self._make_load_post(i)))
 
-    # -- core operations ---------------------------------------------------------
+    # -- state ------------------------------------------------------------------
+    @property
+    def bytes_paged(self) -> int:
+        return sum(self.block_bytes[i] for i in self.paged)
+
+    @property
+    def last(self) -> int:
+        return self._sorted[-1] if self._sorted else -1
+
+    def _resort(self) -> None:
+        self._sorted = sorted(self.paged)
+
+    def _next_paged(self, i: int) -> List[int]:
+        return [j for j in self._sorted if j > i][: self.depth]
+
+    def _prev_paged(self, i: int) -> List[int]:
+        return [j for j in reversed(self._sorted) if j < i][: self.depth]
+
     def _source(self, p, slot) -> torch.Tensor:
         return self.opt.engine.host_param(slot) if slot is not None else self._host_frozen[id(p)]
 
+    # -- moving a block between VRAM and host RAM (between steps) ---------------
+    def page(self, i: int, frozen_init: Optional[Dict[int, torch.Tensor]] = None) -> None:
+        """Make block i paged: its weights leave VRAM."""
+        if i in self.paged or not self._entries[i]:
+            return
+        entries = self._entries[i]
+        slots = [s for _, s, _ in entries if s is not None]
+        if slots:
+            self.opt.engine.wait_ready(slots)  # an in-flight upload into VRAM must land first
+        for p, slot, nbytes in entries:
+            if slot is not None:
+                self.opt.engine.set_paged(slot, True)
+                self.opt._paged_ids.add(id(p))
+                continue
+            # Frozen weight: pinned host copy while that fits the page-lock
+            # budget, else ordinary RAM (the CPU tensor a model was built in is
+            # reused as is). Taken one by one, so a model never exists twice.
+            src = frozen_init.pop(id(p), None) if frozen_init is not None else None
+            if _C.pinned_stats()["live_bytes"] + nbytes <= self.pin_budget:
+                host = _C.pinned_empty([p.numel()], p.dtype, "param.frozen", False, False)
+                host.copy_(src.reshape(-1) if src is not None else p.detach().reshape(-1))
+                self.frozen_pinned_bytes += nbytes
+                host._ts_pinned = True
+            else:
+                host = src.reshape(-1) if src is not None else p.detach().reshape(-1).cpu()
+                self.frozen_pageable_bytes += nbytes
+                host._ts_pinned = False
+            del src
+            self._host_frozen[id(p)] = host
+        self.paged.add(i)
+        self._resort()
+        self.allocated[i] = all(p.untyped_storage().nbytes() > 0 for p, _, _ in entries)
+        self.fresh[i] = False
+        self.release(i)
+        self.stats["pages"] += 1
+
+    def unpage(self, i: int) -> None:
+        """Make block i resident: load its weights and keep them in VRAM."""
+        if i not in self.paged:
+            return
+        self.fresh[i] = False
+        self.load(i)
+        self.wait(i)
+        torch.cuda.current_stream(self.device).synchronize()  # host copies may go only after the copy
+        for p, slot, nbytes in self._entries[i]:
+            if slot is not None:
+                self.opt.engine.set_paged(slot, False)  # updates go straight to VRAM again
+                self.opt._paged_ids.discard(id(p))
+                continue
+            host = self._host_frozen.pop(id(p))
+            if getattr(host, "_ts_pinned", False):
+                self.frozen_pinned_bytes -= nbytes
+            else:
+                self.frozen_pageable_bytes -= nbytes
+            del host
+        self.paged.discard(i)
+        self._resort()
+        self.allocated[i] = True
+        self.stats["unpages"] += 1
+
+    # -- streaming ------------------------------------------------------------------
     def load(self, i: int) -> None:
         """Queue the H2D copy of block i's weights (does not block the GPU)."""
         if i not in self.paged or self.fresh[i]:
@@ -167,31 +237,33 @@ class ParamPager:
         self.stats["releases"] += 1
 
     def materialize(self) -> None:
-        """Load every paged block (e.g. before unwrap)."""
-        for i in self.paged:
-            self.fresh[i] = False
-            self.load(i)
-            self.wait(i)
-        torch.cuda.current_stream(self.device).synchronize()
+        """Make every block resident (e.g. before unwrap)."""
+        for i in list(self._sorted):
+            self.unpage(i)
 
     # -- hooks ---------------------------------------------------------------------
     def _root_pre(self, module, args):
         if not _in_backward():
-            for i in self.paged:  # weights may have been updated since the last load
+            for i in self._sorted:  # weights may have been updated since the last load
                 self.fresh[i] = False
 
     def _make_pre(self, i: int):
         def hook(module, args):
-            self.load(i)
-            self.wait(i)
-            for d in range(1, self.depth + 1):
-                self.load(i + d)
+            if not self._sorted:
+                return
+            if i in self.paged:
+                self.load(i)
+                self.wait(i)
+            for j in self._next_paged(i):
+                self.load(j)
         return hook
 
     def _make_post(self, i: int):
         def hook(module, args, output):
+            if not self._sorted:
+                return
             # autocast keeps a low-precision copy of every weight it casts until the
-            # autocast region exits, which would quietly hold all paged weights in VRAM
+            # autocast region exits, which would quietly hold paged weights in VRAM
             if torch.is_autocast_enabled("cuda"):
                 torch.clear_autocast_cache()
             if _in_backward():
@@ -201,16 +273,17 @@ class ParamPager:
                     if t.requires_grad:
                         t.register_hook(lambda g, i=i: self._on_grad(i))
                 if i == self.last:
-                    return  # first block backward needs
+                    return  # backward starts with this one
             self.release(i)
         return hook
 
     def _on_grad(self, i: int):
-        self.load(i)
-        self.wait(i)
-        for d in range(1, self.depth + 1):
-            self.load(i - d)
-        for j in self.paged:
+        if i in self.paged:
+            self.load(i)
+            self.wait(i)
+        for j in self._prev_paged(i):
+            self.load(j)
+        for j in self._sorted:
             if j > i:
                 self.release(j)
 
@@ -226,6 +299,8 @@ class ParamPager:
 
     def _make_state_dict_hook(self, i: int):
         def hook(module, state_dict, prefix, local_metadata):
+            if i not in self.paged:
+                return
             slots = {id(p): s for p, s, _ in self._entries[i]}
             for name, p in self._named(i):
                 key = prefix + name
@@ -235,12 +310,15 @@ class ParamPager:
 
     def _make_load_pre(self, i: int):
         def hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
-            self.load(i)  # real storage for the copy_ that follows
-            self.wait(i)
+            if i in self.paged:
+                self.load(i)  # real storage for the copy_ that follows
+                self.wait(i)
         return hook
 
     def _make_load_post(self, i: int):
         def hook(module, incompatible_keys):
+            if i not in self.paged:
+                return
             torch.cuda.current_stream(self.device).synchronize()
             for p, slot, _ in self._entries[i]:
                 if slot is not None:
@@ -253,10 +331,6 @@ class ParamPager:
     def remove(self, materialize: bool = True) -> None:
         if materialize:
             self.materialize()
-            for es in self._entries.values():
-                for _, slot, _ in es:
-                    if slot is not None and self.opt.engine is not None:
-                        self.opt.engine.set_paged(slot, False)
         for h in self._handles:
             h.remove()
         self._handles.clear()