block-streamer 0.1.0__py3-none-any.whl

block_streamer/__init__.py

@@ -0,0 +1,5 @@
+"""Public BlockStreamer API."""
+
+from streamer import StreamedModel
+
+__all__ = ["StreamedModel"]

block_streamer/_autograd.py

@@ -0,0 +1,233 @@
+"""Recompute one block at a time; return CPU parameter gradients to autograd."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import torch
+from torch import Tensor
+from torch.autograd.function import once_differentiable
+from torch.utils._pytree import TreeSpec, tree_flatten, tree_unflatten
+
+from memory_pool import tensor_bytes
+
+if TYPE_CHECKING:
+    from streamer import StreamedModel
+
+
+@dataclass
+class TensorTree:
+    """Store structural metadata without retaining live tensor references."""
+
+    leaves: list[Any]
+    positions: list[int]
+    spec: TreeSpec
+
+    @classmethod
+    def split(cls, value: Any) -> tuple[TensorTree, list[Tensor]]:
+        leaves, spec = tree_flatten(value)
+        positions = [i for i, leaf in enumerate(leaves) if isinstance(leaf, Tensor)]
+        tensors = [leaves[i] for i in positions]
+        for i in positions:
+            leaves[i] = None
+        return cls(leaves, positions, spec), tensors
+
+    def merge(self, tensors: list[Tensor] | tuple[Tensor, ...]) -> Any:
+        leaves = self.leaves.copy()
+        for index, tensor in zip(self.positions, tensors, strict=True):
+            leaves[index] = tensor
+        return tree_unflatten(leaves, self.spec)
+
+
+@dataclass
+class Invocation:
+    inputs: TensorTree
+    input_count: int
+    output: TensorTree | None = None
+
+
+class _StreamAutograd(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        model: StreamedModel,
+        call: Invocation,
+        *tensors: Tensor,
+    ) -> tuple[Tensor, ...]:
+        from streamer import hidden_tensor
+
+        ctx.set_materialize_grads(False)
+        ctx.model = model
+        ctx.call = call
+        ctx.configuration_version = model._configuration_version
+        ctx.training_modes = tuple(module.training for module in model.modules())
+        ctx.input_count = call.input_count
+        ctx.master_count = len(tensors) - call.input_count
+        ctx.input_grad_flags = [tensor.requires_grad for tensor in tensors]
+        ctx.autocast = torch.is_autocast_enabled("cuda")
+        ctx.autocast_dtype = torch.get_autocast_dtype("cuda")
+        ctx.rng = []
+        hidden, args, kwargs = call.inputs.merge(tensors[: call.input_count])
+        saved_inputs: list[Tensor] = []
+        result: Any = hidden
+
+        def run(index: int, state: dict[str, Tensor]) -> None:
+            nonlocal result
+            current = hidden_tensor(result)
+            saved_inputs.append(current)
+            ctx.rng.append(
+                (torch.get_rng_state(), torch.cuda.get_rng_state(model.device))
+            )
+            result = model._call_block(index, state, current, args, kwargs)
+
+        model._execute(list(range(len(model.blocks))), run, (hidden, args, kwargs))
+        # Saving CPU master tensors also detects ordinary optimizer/in-place updates
+        # between forward and backward, before stale weights can be re-fetched.
+        ctx.save_for_backward(*tensors, *saved_inputs)
+        call.output, outputs = TensorTree.split(result)
+        ctx.mark_non_differentiable(
+            *[
+                out
+                for out in outputs
+                if not (out.is_floating_point() or out.is_complex())
+            ]
+        )
+        return tuple(outputs)
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx: Any, *output_grads: Tensor | None) -> tuple[Any, ...]:
+        from streamer import hidden_tensor
+
+        model: StreamedModel = ctx.model
+        if model._configuration_version != ctx.configuration_version:
+            raise RuntimeError("Do not call to() between training forward and backward")
+        if tuple(module.training for module in model.modules()) != ctx.training_modes:
+            raise RuntimeError("Do not change train/eval mode before backward")
+        saved = ctx.saved_tensors
+        count = ctx.input_count
+        masters = saved[count : count + ctx.master_count]
+        block_inputs = saved[count + ctx.master_count :]
+        master_indices = {id(tensor): index for index, tensor in enumerate(masters)}
+        input_leaves = [
+            tensor.detach().requires_grad_(ctx.input_grad_flags[index])
+            for index, tensor in enumerate(saved[:count])
+        ]
+        _, args, kwargs = ctx.call.inputs.merge(input_leaves)
+        extra_indices = [
+            index for index in range(1, count) if input_leaves[index].requires_grad
+        ]
+        extra_grads: list[Tensor | None] = [None] * count
+        cpu_contributions: list[tuple[int, Tensor]] = []
+        incoming: Tensor | None = None
+
+        def run(index: int, state: dict[str, Tensor]) -> None:
+            nonlocal incoming
+            with (
+                torch.enable_grad(),
+                torch.random.fork_rng(devices=[model.device.index]),
+            ):
+                torch.set_rng_state(ctx.rng[index][0])
+                torch.cuda.set_rng_state(ctx.rng[index][1], model.device)
+                hidden = block_inputs[index].detach().requires_grad_(True)
+                names = [
+                    name
+                    for name, parameter in model.blocks[index].named_parameters()
+                    if parameter.requires_grad
+                ]
+                targets = (
+                    [hidden]
+                    + [input_leaves[i] for i in extra_indices]
+                    + [state[name] for name in names]
+                )
+                with torch.autocast(
+                    "cuda",
+                    enabled=ctx.autocast,
+                    dtype=ctx.autocast_dtype,
+                    # An outer autocast scope must not cache converted weights
+                    # after this block retires, defeating parameter residency.
+                    cache_enabled=False,
+                ):
+                    result = model._call_block(index, state, hidden, args, kwargs)
+                if index == len(model.blocks) - 1:
+                    _, outputs = TensorTree.split(result)
+                    pairs = [
+                        (out, grad)
+                        for out, grad in zip(outputs, output_grads, strict=True)
+                        if out.requires_grad and grad is not None
+                    ]
+                else:
+                    out = hidden_tensor(result)
+                    pairs = (
+                        [(out, incoming)]
+                        if out.requires_grad and incoming is not None
+                        else []
+                    )
+                if pairs:
+                    gradients = torch.autograd.grad(
+                        [pair[0] for pair in pairs],
+                        targets,
+                        [pair[1] for pair in pairs],
+                        allow_unused=True,
+                    )
+                else:
+                    gradients = (None,) * len(targets)
+            incoming = gradients[0]
+            for input_index, gradient in zip(
+                extra_indices, gradients[1:], strict=False
+            ):
+                if gradient is not None:
+                    previous = extra_grads[input_index]
+                    extra_grads[input_index] = (
+                        gradient if previous is None else previous + gradient
+                    )
+            parameters = dict(model.blocks[index].named_parameters())
+            for name, gradient in zip(
+                names, gradients[1 + len(extra_indices) :], strict=True
+            ):
+                if gradient is None:
+                    continue
+                parameter = parameters[name]
+                cpu = torch.empty_like(parameter, device="cpu", pin_memory=True)
+                # D2H runs on compute after the gradient kernels. Slot retirement's
+                # completion event also proves these pinned destinations are readable.
+                cpu.copy_(gradient, non_blocking=True)
+                gradient.record_stream(model.compute_stream)
+                model.gradient_transfer_bytes += tensor_bytes(gradient)
+                cpu_contributions.append((master_indices[id(parameter)], cpu))
+
+        model._execute(
+            list(reversed(range(len(model.blocks)))),
+            run,
+            (block_inputs, input_leaves, output_grads),
+            grad=True,
+        )
+        # _execute has retired all completion events, including the D2H copies.
+        master_grads: list[Tensor | None] = [None] * len(masters)
+        for index, gradient in cpu_contributions:
+            if master_grads[index] is None:
+                master_grads[index] = gradient
+            else:
+                master_grads[index].add_(gradient)
+        extra_grads[0] = incoming if ctx.input_grad_flags[0] else None
+        return (None, None, *extra_grads, *master_grads)
+
+
+def training_forward(
+    model: StreamedModel,
+    hidden: Tensor,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> Any:
+    if not model.blocks:
+        return hidden
+    if not hidden.is_floating_point():
+        raise TypeError("Training requires floating-point hidden states")
+    tree, inputs = TensorTree.split((hidden, args, kwargs))
+    call = Invocation(tree, len(inputs))
+    # Include buffers in saved state so registered-state changes are version checked.
+    masters = list(model.parameters()) + list(model.buffers())
+    outputs = _StreamAutograd.apply(model, call, *inputs, *masters)
+    assert call.output is not None
+    return call.output.merge(outputs)

block_streamer-0.1.0.dist-info/METADATA

@@ -0,0 +1,298 @@
+Metadata-Version: 2.4
+Name: block-streamer
+Version: 0.1.0
+Summary: Bounded, asynchronous parameter streaming for PyTorch blocks
+License-Expression: Apache-2.0
+Project-URL: Repository, https://github.com/hareramray/BlockStreamer
+Project-URL: Issues, https://github.com/hareramray/BlockStreamer/issues
+Keywords: pytorch,cuda,offloading,inference,training
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.4
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: ruff>=0.9; extra == "dev"
+Requires-Dist: matplotlib>=3.8; extra == "dev"
+Dynamic: license-file
+
+# BlockStreamer
+
+Licensed under the [Apache License 2.0](https://github.com/hareramray/BlockStreamer/blob/main/LICENSE).
+
+Execute ordered PyTorch blocks with pinned CPU master weights, asynchronous CUDA
+prefetch, and at most `prefetch_ahead + 1` resident blocks. Includes inference,
+first-order training with backward re-fetch, correctness tests, and a measured
+transfer/compute crossover benchmark.
+
+## Prior art and binding decision
+
+Accelerate's `cpu_offload_with_hook` leaves a module on the execution device until
+its offload hook runs; BlockStreamer instead schedules an explicit bounded window
+of individual blocks and transfer-ready events.
+[Accelerate documentation](https://huggingface.co/docs/accelerate/en/package_reference/big_modeling#accelerate.cpu_offload_with_hook)
+ZeRO-Infinity is a broader training system that spans GPU, CPU, and NVMe memory;
+this project targets one GPU, CPU parameter storage, and an ordered block list.
+[ZeRO-Infinity paper](https://arxiv.org/abs/2104.07857)
+FSDP `CPUOffload` moves parameters and gradients to CPU within FSDP's sharding and
+distributed training machinery; this library has no distributed collectives or
+parameter sharding.
+[FSDP documentation](https://docs.pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.CPUOffload)
+BlockStreamer's smaller scope makes the residency policy explicit, at the cost of
+less general model execution and substantial Python/event overhead for tiny blocks.
+
+**Binding mechanism: `torch.func.functional_call`.** Each block executes with a
+dictionary of its staged GPU parameters and buffers, with strict key checking and
+PyTorch's tied-weight handling. Registered parameters stay on pinned CPU storage.
+Construction and `.to()` replace CPU storage to pin/convert it while preserving
+`Parameter` identity; `.data` is never used to bind staged GPU weights for compute.
+[Functional-call documentation](https://docs.pytorch.org/docs/stable/generated/torch.func.functional_call.html)
+
+The tradeoff is functional state discipline: blocks must not mutate registered
+parameters or buffers, retain staged tensors, or use custom kernels that cache
+parameter pointers outside module lookups. Detected registered-state mutation
+raises an error. Python attributes, global state, hooks, and opaque external caches
+cannot all be audited automatically; callers must keep those free of side effects.
+Training BatchNorm's running-state updates and mutable KV caches are unsupported.
+Ordinary dropout is supported through RNG replay during recomputation.
+
+## Install and run
+
+Requires Python 3.11+, PyTorch 2.4+, and a CUDA GPU. Install a CUDA-enabled PyTorch
+wheel suitable for your GPU first. For the published package:
+
+```bash
+python -m pip install block-streamer
+```
+
+For development, clone the repository and install from its root:
+
+```bash
+python -m pip install -e ".[dev]"
+python -m pytest -q
+```
+
+The development environment here uses `.venv` with the preinstalled CUDA PyTorch
+available through system site packages. On this Windows workspace:
+
+```powershell
+.venv/Scripts/python.exe -m pytest -q
+.venv/Scripts/python.exe benchmark.py --widths 128 256 512 1024 --tokens 32 2048 8192 --runs 10 --warmup 4
+```
+
+```python
+import torch
+from torch import nn
+from block_streamer import StreamedModel
+
+blocks = nn.Sequential(nn.Linear(256, 512), nn.GELU(), nn.Linear(512, 256))
+x = torch.randn(16, 256, device="cuda:0", dtype=torch.bfloat16)
+
+with StreamedModel(
+    blocks, device="cuda:0", prefetch_ahead=1, dtype=torch.bfloat16
+) as streamed:
+    out = streamed(x)
+    print(streamed.stats())
+```
+
+Construction takes ownership of the supplied `Sequential`, `ModuleList`, or list
+of modules, pins their parameters **and buffers** on CPU, and defaults to `.eval()`.
+Make a deepcopy first if another model must retain independent state. Evaluation
+forwards disable autograd, including when the caller did not use `no_grad()`.
+Call `.train()` to enable the training path. `.eval()` propagates to every block.
+
+Inputs must already be on the execution device and have a suitable dtype. Additional
+positional and keyword arguments are passed unchanged to **every** block:
+
+```python
+out = streamed(x, attention_mask=mask, position_ids=positions)
+```
+
+Every block must accept the supplied arguments. A block may return a tensor or a
+tuple whose first element is the next hidden state. Intermediate auxiliary tuple
+elements are discarded; the final block's entire result is returned. Tensor leaves
+in ordinary tuple/list/dict output trees are cloned so a returned parameter view
+cannot keep staged weights resident. There is no automatic per-layer KV-cache
+selection or collection; use adapter blocks to express that policy. Read-only
+attention masks, positions, rotary tensors, and cache values can pass through.
+
+`streamed.to(device="cuda:0", dtype=torch.float16)` changes execution placement and
+repins CPU state. Perform conversion before constructing an optimizer or a pending
+training graph. CPU execution and inherited `.cuda()`, `.cpu()`, `.half()`, and
+parent-module conversions are rejected; use this wrapper's `.to()` instead.
+`close()` and context exit drain outstanding work; the wrapper remains reusable.
+
+## Streams, residency, and memory budget
+
+The engine uses two streams: the compute stream active at construction and a
+separate transfer stream. Borrowing the compute stream avoids creating persistent
+cuBLAS workspaces for every new wrapper. A call made on another CUDA stream is
+bridged to the compute stream, with input storage protected by `record_stream()`.
+
+The scheduler initially stages up to `prefetch_ahead + 1` blocks. Compute waits on
+each block's H2D-ready event. **Every transfer-created tensor consumed by compute
+calls `record_stream(compute_stream)` in `BufferPool.acquire()`**, protecting it
+from premature caching-allocator reuse.
+[PyTorch allocator documentation](https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html)
+A compute-completion event is host-synchronized before eviction and replacement
+staging. This conservative retirement rule enforces the actual live-block limit,
+including in-flight transfers, and introduces host scheduling overhead. The forward
+returns after its work completes; H2D and compute can still overlap within it.
+
+The GPU pool uses PyTorch's caching allocator with event-retired slots. Allocations
+are sized per block, so heterogeneous shapes, integer buffers, and mixed state
+dtypes require no homogeneous-buffer fallback. `prefetch_ahead=0` is supported;
+`1` provides two slots. Parameter-free blocks still count as logical slots.
+CUDA **reserved** cache memory may exceed live tensor payload and is not the
+residency invariant. For heterogeneous models the byte bound is the largest live
+window's summed state size, not a fixed number of bytes per block.
+
+Set `buffer_budget_bytes=...` for a per-slot parameter-plus-buffer payload limit.
+An oversized block raises a clear error before its copy. This is not a total-VRAM
+budget: activations, CUDA library workspaces, gradients, and allocator rounding
+are additional. Every H2D source is checked for pinned CPU storage immediately
+before copying; an unpinned source fails loudly. Meta/sparse state and distinct
+tensor objects sharing backing storage are rejected. Weight tying by sharing the
+same `Parameter` object is supported within and across blocks.
+
+`stats()` reports cumulative H2D bytes/time, D2H gradient bytes, GPU transfer-wait
+time (`stall_ms`), host eviction wait time, current/peak resident blocks, and peak
+parameter/state payload bytes. GPU stall time measures event waits on the compute
+stream; it excludes host launch delays. Host eviction wait can include compute
+and transfer waiting, so these two times must not be added as independent costs.
+`residency_history` retains the last 4,096 transitions; an optional
+`residency_observer` receives every transition. `reset_stats()` resets all counters.
+
+## Training design and usage
+
+The inference suite passed on CUDA before the training implementation was added.
+The training design makes these choices:
+
+1. **Re-fetch weights during backward.** Keeping forward weights would require
+   O(all blocks) parameter VRAM. Instead, save block inputs on GPU, discard staged
+   weights, and re-fetch/recompute one block at a time during backward. A training
+   step transfers each block's state twice and performs an extra forward compute.
+   Saved activations remain O(depth); activation offloading is outside this project.
+2. **Offload gradients after each block.** Compute that block's gradients on GPU,
+   copy them into pinned CPU tensors on the compute stream, and make CPU reads wait
+   for the completion event. Return CPU parameter gradients through autograd so
+   normal `.grad` accumulation works. Peak GPU memory also includes one block's
+   gradients and recomputation graph, plus saved activations and input gradients.
+   Shared parameters' CPU gradient contributions are summed after copies finish.
+3. **CPU optimizer and state.** Use an ordinary CPU optimizer on `model.parameters()`.
+   The optimizer step runs on CPU and does not stream to GPU. This incurs host RAM
+   for all gradients/optimizer state. No automatic fp32 master-weight conversion,
+   fused offloaded optimizer, or integrated loss scaling is provided.
+4. **Custom `torch.autograd.Function`.** It owns the forward/backward schedule,
+   saves version-checked master state and block inputs, then visits blocks in
+   reverse with the same bounded prefetch window. Block N-1 is fetched ahead of
+   its recompute/backward while block N executes. This avoids relying on module
+   hook order and explicitly returns gradients for tensor arguments and weights.
+
+```python
+streamed = StreamedModel(blocks, dtype=torch.float32).train()
+optimizer = torch.optim.Adam(streamed.parameters(), lr=1e-3)
+x = torch.randn(16, 256, device="cuda")
+target = torch.zeros_like(x)
+
+for _ in range(10):
+    optimizer.zero_grad(set_to_none=True)
+    loss = (streamed(x) - target).float().square().mean()
+    loss.backward()
+    optimizer.step()  # CPU parameters, CPU gradients, CPU optimizer state
+```
+
+Recomputation restores PyTorch CPU/CUDA RNG states and CUDA autocast settings.
+Backward does not advance the caller's RNG. First-order gradients, frozen/unused
+parameters, tied weights, extra tensor arguments, and final tuple-output losses
+are supported. Higher-order differentiation, `torch.compile`, CUDA graph capture,
+distributed execution, and concurrent/reentrant calls are unsupported. Modules
+must not modify block inputs in place, change execution behavior between forward
+and backward, or retain GPU weights in external state. Normal in-place changes to
+saved tensors are version-checked; edits through `.data` bypass PyTorch's protection.
+
+## Measured crossover and memory
+
+![Measured transfer/compute crossover and peak CUDA allocation](https://raw.githubusercontent.com/hareramray/BlockStreamer/main/results/roofline.png)
+
+Local measurement: NVIDIA GeForce RTX 5050 Laptop GPU (8 GB), Windows WDDM,
+PyTorch 2.11.0+cu128 / CUDA runtime 12.8, bf16, six heterogeneous MLP blocks,
+`prefetch_ahead=1`, four warmups and ten timed runs. These are measurements on this
+machine, not predictions for a desktop PCIe Gen4 x16 link. The benchmark records
+31.5 GB/s as the nominal Gen4 x16 reference but uses **measured pinned H2D time**
+for its conclusions. See [the complete report](https://github.com/hareramray/BlockStreamer/blob/main/results/benchmark.json).
+
+**Headline:** at 32 and 2,048 tokens, the sampled compute/H2D crossover falls
+between widths 128 and 256. At 8,192 tokens, compute time exceeds H2D time at every
+sampled width; no width crossover was observed. At width 1,024, increasing tokens
+from 2,048 to 8,192 changes the measured regime from bandwidth-bound to compute-bound.
+Small-block CUDA event intervals include host launch gaps; the small-width
+"compute-bound" label means the measured compute path takes longer than H2D,
+not that GPU arithmetic units are saturated. This sweep is an empirical roofline
+diagnostic, not a hardware FLOP/s roofline or a universal crossover constant.
+
+| Width / tokens | GPU median ms | Streamed median ms | GPU peak MiB | Streamed peak MiB | Activation/workspace increment MiB |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| 1,024 / 32 | 0.569 | 10.524 | 68.625 | 28.625 | 0.438 |
+| 1,024 / 2,048 | 5.915 | 14.883 | 100.125 | 60.125 | 28.000 |
+| 1,024 / 8,192 | 24.304 | 29.825 | 196.125 | 156.125 | 112.000 |
+
+For these width-1,024 cases, baseline parameter payload is 60 MiB and streamed
+peak parameter payload is 20 MiB. The 8,192-token activation/workspace increment
+alone is 112 MiB: parameter streaming does not eliminate activation-dominated VRAM.
+This implementation saves memory but **was slower in every sampled case**.
+
+The benchmark reports raw `torch.cuda.max_memory_allocated()` for total peak VRAM,
+exact parameter/state tensor payload maxima, and an **independently measured**
+activation/workspace increment: peak baseline allocation above the warmed baseline
+with weights already resident. Existing input/library allocations are also recorded.
+Separate maxima occur at different times; subtracting peak parameter payload from
+total peak does not produce an exact activation peak. PyTorch's aggregate allocator
+counter also cannot distinguish an activation from a temporary workspace.
+Therefore an exact additive parameter-versus-activation split is not claimed.
+
+Latency uses synchronized wall-clock intervals and reports medians, means,
+population variance, minima, and maxima. The requested overlap score is:
+
+```text
+1 - (streamed_median_ms - all_on_GPU_median_ms) / isolated_H2D_ms
+```
+
+It is not clamped: Python, event, cloning, and scheduling overhead can make it
+negative; noise can put it above one. The score is 0.454 for width 1,024 / 8,192
+tokens in this run. Kernel overlap cannot fully hide transfer when per-block H2D
+time exceeds compute time. Increasing token count changes arithmetic intensity;
+increasing width alone need not produce a crossover.
+
+## Profiler and validation status
+
+`benchmark.py` exports `results/trace.json`, readable by Chrome tracing or Perfetto,
+and checks whether GPU activities are present. On this machine CUPTI initialization
+returned `CUPTI_ERROR_INVALID_DEVICE`, so the exported trace contains **CPU events
+only** and `trace_gpu_activity` is `false`. GPU timeline overlap is therefore **not
+visually verified here**. CUDA-event timing, real GPU parity tests, and memory
+measurements did run. Re-run on a compatible CUPTI/driver setup to capture GPU
+kernels and copies; `--no-trace` explicitly skips profiler collection.
+
+The final local suite passed **32 tests** on CUDA, with Ruff lint/format checks
+also passing. Tests cover inference in fp32/bf16/fp16 at prefetch depths 0–3,
+heterogeneous blocks and nonpersistent buffers, extra arguments and tuple results,
+every residency transition, a 100-iteration bandwidth-starved allocator stress,
+nondefault caller streams, inference mode, output aliases, tied weights, unpinned
+and over-budget errors, and exception cleanup. Training tests cover fp32/bf16/fp16
+gradient parity at depths 1–3, CPU pinned gradients, extra-input gradients, dropout
+RNG replay, unused/shared parameters, accumulation, autocast with frozen weights,
+CPU Adam state placement, and a 12-step SGD toy loss curve.
+
+Inference `atol=rtol` is `1e-5` for fp32, `8e-3` for bf16, and `1e-3` for fp16,
+allowing roughly one low-precision rounding unit at unit scale. Backward
+`(atol, rtol)` is `(2e-5, 2e-4)`, `(2e-3, 3e-2)`, and `(5e-4, 5e-3)` respectively;
+backward includes reductions and multiple gradient contributions. The convergence
+curves use fp32 with `(2e-6, 2e-5)` tolerances and must decrease by at least 10%.
+These checks are evidence for the tested workloads, not proof for arbitrary custom
+modules or hardware. The environment emitted a Triton CUDA-toolkit discovery
+warning; the tests use PyTorch CUDA operations and do not require Triton compilation.

block_streamer-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+memory_pool.py,sha256=0WZRxiW40QskJyozR3sLYM605Xav8krExLcb7SGKJU4,7787
+streamer.py,sha256=Z6q3lXutVqRHTXPxgeGu32je-bJxS89L69HnUqb6Bb0,9387
+block_streamer/__init__.py,sha256=oyBCIsDinH6UddoTM5dcUcSXUWm-Mz5A-f_7sUcbvfA,97
+block_streamer/_autograd.py,sha256=P5N6PrNFfKq5aqywdqpCGinfxHM57HKt2lWaBjYmQbI,9438
+block_streamer-0.1.0.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+block_streamer-0.1.0.dist-info/METADATA,sha256=0ApH7YgRl9A4eg6J3EJAUdbZniVNBB5wLfbAtsvUg_U,17177
+block_streamer-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+block_streamer-0.1.0.dist-info/top_level.txt,sha256=BpKpp9UuLaPkcBAweNHk_qJ0OpO37rcLzlI5jTofnKM,36
+block_streamer-0.1.0.dist-info/RECORD,,

block_streamer-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

block_streamer-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

block_streamer-0.1.0.dist-info/top_level.txt

@@ -0,0 +1,3 @@
+block_streamer
+memory_pool
+streamer

memory_pool.py

@@ -0,0 +1,191 @@
+"""Pinned canonical state and bounded, event-retired CUDA allocations."""
+
+from __future__ import annotations
+
+import time
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+
+
+def tensor_bytes(tensor: Tensor) -> int:
+    return tensor.numel() * tensor.element_size()
+
+
+def module_state(module: nn.Module) -> dict[str, Tensor]:
+    """Include nonpersistent buffers; functional_call needs those too."""
+    return dict(list(module.named_parameters()) + list(module.named_buffers()))
+
+
+def pin_modules(blocks: nn.ModuleList, dtype: torch.dtype | None) -> None:
+    """Preserve Parameter identity and shared objects when creating CPU storage."""
+    if dtype is not None and not dtype.is_floating_point:
+        raise TypeError("dtype must be a floating-point torch.dtype")
+    storage_owners: dict[tuple[torch.device, int], int] = {}
+    tensors = list(blocks.parameters()) + list(blocks.buffers())
+    for tensor in tensors:
+        if tensor.is_meta or tensor.layout != torch.strided:
+            raise ValueError("Only materialized, dense strided state is supported")
+        if tensor.numel():
+            key = (tensor.device, tensor.untyped_storage().data_ptr())
+            if key in storage_owners and storage_owners[key] != id(tensor):
+                raise ValueError(
+                    "Distinct state tensors sharing storage are unsupported; "
+                    "tie weights by sharing the same Parameter object instead"
+                )
+            storage_owners[key] = id(tensor)
+    seen: set[int] = set()
+    for tensor in tensors:
+        if id(tensor) in seen:
+            continue
+        seen.add(id(tensor))
+        if tensor.is_meta or tensor.layout != torch.strided:
+            raise ValueError("Only materialized, dense strided state is supported")
+        target_dtype = dtype if tensor.is_floating_point() else None
+        cpu = tensor.detach().to(device="cpu", dtype=target_dtype).pin_memory()
+        tensor.data = cpu
+        if isinstance(tensor, nn.Parameter) and tensor.grad is not None:
+            tensor.grad = tensor.grad.to(device="cpu", dtype=cpu.dtype).pin_memory()
+
+
+@dataclass
+class Resident:
+    state: dict[str, Tensor]
+    ready: torch.cuda.Event
+    start: torch.cuda.Event
+    nbytes: int
+    parameter_bytes: int
+
+
+class BufferPool:
+    """Each live slot owns one block, regardless of shape or dtype.
+
+    Allocations use PyTorch's caching allocator rather than shape-specific slabs.
+    A slot is freed only after its compute completion event has completed. Cached
+    (reserved) allocator memory is not live parameter residency.
+    """
+
+    def __init__(
+        self,
+        device: torch.device,
+        capacity: int,
+        transfer_stream: torch.cuda.Stream,
+        compute_stream: torch.cuda.Stream,
+        budget_bytes: int | None = None,
+        observer: Callable[[int], None] | None = None,
+    ) -> None:
+        self.device = device
+        self.capacity = capacity
+        self.transfer_stream = transfer_stream
+        self.compute_stream = compute_stream
+        self.budget_bytes = budget_bytes
+        self.observer = observer
+        self.live: dict[int, Resident] = {}
+        self.reset_stats()
+
+    def reset_stats(self) -> None:
+        if self.live:
+            raise RuntimeError("Cannot reset statistics while blocks are resident")
+        self.transfer_bytes = 0
+        self.transfer_ms = 0.0
+        self.stall_ms = 0.0
+        self.eviction_wait_ms = 0.0
+        self.peak_resident_blocks = 0
+        self.peak_state_bytes = 0
+        self.peak_parameter_bytes = 0
+        self.history: deque[int] = deque(maxlen=4096)
+        self.wait_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = []
+
+    def check(self) -> None:
+        count = len(self.live)
+        assert count <= self.capacity, "GPU block residency exceeded capacity"
+        self.history.append(count)
+        self.peak_resident_blocks = max(self.peak_resident_blocks, count)
+        self.peak_state_bytes = max(
+            self.peak_state_bytes, sum(item.nbytes for item in self.live.values())
+        )
+        self.peak_parameter_bytes = max(
+            self.peak_parameter_bytes,
+            sum(item.parameter_bytes for item in self.live.values()),
+        )
+        if self.observer is not None:
+            self.observer(count)
+
+    def stage(self, index: int, module: nn.Module, grad: bool = False) -> None:
+        if index in self.live:
+            return
+        if len(self.live) >= self.capacity:
+            raise RuntimeError("No free GPU block slot; retire completed compute first")
+        source = module_state(module)
+        for name, tensor in source.items():
+            if tensor.device.type != "cpu" or not tensor.is_pinned():
+                raise RuntimeError(
+                    f"Block {index} tensor {name!r} is not pinned CPU memory"
+                )
+        nbytes = sum(tensor_bytes(tensor) for tensor in source.values())
+        if self.budget_bytes is not None and nbytes > self.budget_bytes:
+            raise ValueError(
+                f"Block {index} requires {nbytes} bytes, exceeding per-slot buffer "
+                f"budget {self.budget_bytes} bytes"
+            )
+        parameters = dict(module.named_parameters())
+        with torch.cuda.stream(self.transfer_stream):
+            start = torch.cuda.Event(enable_timing=True)
+            ready = torch.cuda.Event(enable_timing=True)
+            start.record(self.transfer_stream)
+            state = {
+                name: tensor.detach().to(self.device, non_blocking=True)
+                for name, tensor in source.items()
+            }
+            if grad:
+                for name, parameter in parameters.items():
+                    state[name].requires_grad_(parameter.requires_grad)
+            ready.record(self.transfer_stream)
+        self.live[index] = Resident(
+            state,
+            ready,
+            start,
+            nbytes,
+            sum(tensor_bytes(tensor) for tensor in parameters.values()),
+        )
+        self.transfer_bytes += nbytes
+        self.check()
+
+    def acquire(self, index: int) -> dict[str, Tensor]:
+        item = self.live[index]
+        before = torch.cuda.Event(enable_timing=True)
+        after = torch.cuda.Event(enable_timing=True)
+        before.record(self.compute_stream)
+        # Compute must not read weights until their H2D copies have finished.
+        self.compute_stream.wait_event(item.ready)
+        after.record(self.compute_stream)
+        self.wait_events.append((before, after))
+        for tensor in item.state.values():
+            # CRITICAL allocator guard: transfer-created storage remains protected
+            # while compute consumes it, even if Python drops its final reference.
+            tensor.record_stream(self.compute_stream)
+        return item.state
+
+    def evict(self, index: int, done: torch.cuda.Event) -> None:
+        start = time.perf_counter()
+        # Host retirement proves compute is complete BEFORE allocating a new slot;
+        # merely enqueueing a wait would permit too many live parameter blocks.
+        done.synchronize()
+        self.eviction_wait_ms += (time.perf_counter() - start) * 1000
+        item = self.live.pop(index)
+        self.transfer_ms += item.start.elapsed_time(item.ready)
+        item.state.clear()
+        self.check()
+
+    def drain(self) -> None:
+        # On exceptions, both pending copies and consumers must finish before free.
+        self.transfer_stream.synchronize()
+        self.compute_stream.synchronize()
+        self.live.clear()
+        self.check()
+        for before, after in self.wait_events:
+            self.stall_ms += before.elapsed_time(after)
+        self.wait_events.clear()

streamer.py

@@ -0,0 +1,234 @@
+"""Blockwise functional execution with bounded CUDA parameter residency."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import Any, Self
+
+import torch
+from torch import Tensor, nn
+from torch.func import functional_call
+from torch.utils._pytree import tree_map
+
+from memory_pool import BufferPool, pin_modules
+
+
+def hidden_tensor(result: Any) -> Tensor:
+    hidden = result[0] if isinstance(result, tuple) else result
+    if not isinstance(hidden, Tensor):
+        raise TypeError(
+            "A block must return a Tensor or a tuple beginning with a Tensor"
+        )
+    return hidden
+
+
+class StreamedModel(nn.Module):
+    """Sequential block streaming. Intermediate tuple[0] feeds the next block.
+
+    Additional positional/keyword arguments are supplied unchanged to every block;
+    the final block's entire result is returned. Inputs must already be on device.
+    Construction takes ownership of the supplied modules and pins their CPU state.
+    """
+
+    def __init__(
+        self,
+        blocks: nn.Sequential | nn.ModuleList | Sequence[nn.Module],
+        device: str | torch.device = "cuda:0",
+        prefetch_ahead: int = 1,
+        dtype: torch.dtype | None = None,
+        buffer_budget_bytes: int | None = None,
+        residency_observer: Callable[[int], None] | None = None,
+    ) -> None:
+        super().__init__()
+        if isinstance(prefetch_ahead, bool) or not isinstance(prefetch_ahead, int):
+            raise TypeError("prefetch_ahead must be an integer")
+        if prefetch_ahead < 0:
+            raise ValueError("prefetch_ahead must be nonnegative")
+        if buffer_budget_bytes is not None and buffer_budget_bytes <= 0:
+            raise ValueError("buffer_budget_bytes must be positive")
+        self.device = self._cuda_device(device)
+        self.blocks = nn.ModuleList(list(blocks))
+        self.prefetch_ahead = prefetch_ahead
+        self.dtype = dtype
+        self.buffer_budget_bytes = buffer_budget_bytes
+        self.residency_observer = residency_observer
+        self._busy = False
+        self._configuration_version = 0
+        self.gradient_transfer_bytes = 0
+        pin_modules(self.blocks, dtype)
+        self._make_engine()
+        self.eval()
+
+    @staticmethod
+    def _cuda_device(device: str | torch.device) -> torch.device:
+        result = torch.device(device)
+        if result.type != "cuda":
+            raise ValueError("StreamedModel requires a CUDA execution device")
+        if not torch.cuda.is_available():
+            raise RuntimeError(
+                "CUDA is unavailable; install a CUDA-enabled PyTorch build"
+            )
+        index = torch.cuda.current_device() if result.index is None else result.index
+        if not 0 <= index < torch.cuda.device_count():
+            raise ValueError(f"CUDA device index {index} is unavailable")
+        return torch.device("cuda", index)
+
+    def _make_engine(self) -> None:
+        # Borrow the construction stream. A new compute stream per wrapper would
+        # create a persistent cuBLAS workspace per instance in PyTorch's handle pool.
+        self.compute_stream = torch.cuda.current_stream(device=self.device)
+        self.transfer_stream = torch.cuda.Stream(device=self.device)
+        self.pool = BufferPool(
+            self.device,
+            self.prefetch_ahead + 1,
+            self.transfer_stream,
+            self.compute_stream,
+            self.buffer_budget_bytes,
+            self.residency_observer,
+        )
+
+    def _execute(
+        self,
+        order: list[int],
+        callback: Callable[[int, dict[str, Tensor]], None],
+        inputs: Any,
+        grad: bool = False,
+    ) -> None:
+        if self._busy:
+            raise RuntimeError("Concurrent or reentrant execution is unsupported")
+        self._busy = True
+        caller = torch.cuda.current_stream(self.device)
+        # Inputs may have been produced on the caller stream just before forward.
+        self.compute_stream.wait_stream(caller)
+
+        def protect(value: Any) -> Any:
+            if isinstance(value, Tensor) and value.is_cuda:
+                if value.device != self.device:
+                    raise ValueError(
+                        f"Input is on {value.device}, expected {self.device}"
+                    )
+                value.record_stream(self.compute_stream)
+            return value
+
+        try:
+            tree_map(protect, inputs)
+            for index in order[: self.pool.capacity]:
+                self.pool.stage(index, self.blocks[index], grad)
+            for offset, index in enumerate(order):
+                with torch.cuda.stream(self.compute_stream):
+                    callback(index, self.pool.acquire(index))
+                    done = torch.cuda.Event()
+                    done.record(self.compute_stream)
+                self.pool.evict(index, done)
+                next_offset = offset + self.pool.capacity
+                if next_offset < len(order):
+                    upcoming = order[next_offset]
+                    self.pool.stage(upcoming, self.blocks[upcoming], grad)
+        finally:
+            self.pool.drain()
+            # Returning tensors to caller must not race their production on compute.
+            caller.wait_stream(self.compute_stream)
+            self._busy = False
+
+    def _call_block(
+        self,
+        index: int,
+        state: dict[str, Tensor],
+        hidden: Tensor,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        versions = {name: value._version for name, value in state.items()}
+        identities = {name: id(value) for name, value in state.items()}
+        result = functional_call(
+            self.blocks[index],
+            state,
+            (hidden, *args),
+            kwargs,
+            strict=True,
+        )
+        if any(
+            value._version != versions[name] or id(value) != identities[name]
+            for name, value in state.items()
+        ):
+            raise RuntimeError(
+                "Blocks must not mutate parameters or registered buffers"
+            )
+        hidden_tensor(result)
+        # A block may return a view of a weight. Clone outputs so eviction really
+        # releases the staged state and callers cannot retain the pool's storage.
+        return tree_map(
+            lambda value: value.clone() if isinstance(value, Tensor) else value, result
+        )
+
+    def forward(self, hidden_states: Tensor, *args: Any, **kwargs: Any) -> Any:
+        if hidden_states.device != self.device:
+            raise ValueError(f"hidden_states must be on {self.device}")
+        if self.training and torch.is_grad_enabled():
+            from block_streamer._autograd import training_forward
+
+            return training_forward(self, hidden_states, args, kwargs)
+        result: Any = hidden_states
+
+        def run(index: int, state: dict[str, Tensor]) -> None:
+            nonlocal result
+            result = self._call_block(index, state, hidden_tensor(result), args, kwargs)
+
+        with torch.inference_mode(False), torch.no_grad():
+            self._execute(
+                list(range(len(self.blocks))), run, (hidden_states, args, kwargs)
+            )
+        return result
+
+    def to(self, *args: Any, **kwargs: Any) -> Self:
+        """Change execution device/dtype while retaining pinned CPU master state."""
+        if self._busy:
+            raise RuntimeError("Cannot move the model during execution")
+        device, dtype, _, memory_format = torch._C._nn._parse_to(*args, **kwargs)
+        if memory_format is not None:
+            raise ValueError("memory_format conversion is unsupported")
+        target = self.device if device is None else self._cuda_device(device)
+        self.close()
+        self._configuration_version += 1
+        pin_modules(self.blocks, dtype)
+        self.device = target
+        if dtype is not None:
+            self.dtype = dtype
+        self._make_engine()
+        return self
+
+    def _apply(self, fn: Callable[[Tensor], Tensor], recurse: bool = True) -> Self:
+        raise RuntimeError(
+            "Use StreamedModel.to(device=..., dtype=...) to preserve pinned CPU state; "
+            "inherited cpu/cuda/half/bfloat16 and parent-module conversions are unsupported"
+        )
+
+    def stats(self) -> dict[str, Any]:
+        """Cumulative since construction/reset; byte counts are tensor payloads."""
+        return {
+            "transfer_bytes": self.pool.transfer_bytes,
+            "gradient_transfer_bytes": self.gradient_transfer_bytes,
+            "transfer_ms": self.pool.transfer_ms,
+            "stall_ms": self.pool.stall_ms,
+            "eviction_wait_ms": self.pool.eviction_wait_ms,
+            "resident_blocks": len(self.pool.live),
+            "peak_resident_blocks": self.pool.peak_resident_blocks,
+            "peak_state_bytes": self.pool.peak_state_bytes,
+            "peak_parameter_bytes": self.pool.peak_parameter_bytes,
+            "residency_history": tuple(self.pool.history),
+        }
+
+    def reset_stats(self) -> None:
+        self.pool.reset_stats()
+        self.gradient_transfer_bytes = 0
+
+    def close(self) -> None:
+        if self._busy:
+            raise RuntimeError("Cannot close the model during execution")
+        self.pool.drain()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()