Piper Offload

A model-agnostic GPU/CPU memory manager for PyTorch. It caches reusable model and LoRA resources, and swaps independent models in and out of GPU memory under a policy-driven cache.

Piper Offload is self-contained and library-friendly: it has no required dependency beyond torch. Optional integrations support bitsandbytes, optimum.quanto, gguf, piper-kernels, and torchao quantized models.

Requires Python 3.14 or newer and PyTorch 2.13.

Installation

Install the base package from PyPI:

pip install piper-offload

Optional integrations are available individually through the bnb, torchao, gguf, quanto, and convrot extras. Triton acceleration is a separate, composable extra, while all includes every integration plus Triton:

pip install "piper-offload[all]"

The triton extra selects upstream triton on Linux and triton-windows on 64-bit Windows. Combine it with any individual quantization extra whose optimized kernels you want; without it, those integrations retain their portable fallback paths. The all extra includes this acceleration runtime. The same installed Triton runtime enables Piper Kernels' ConvRot backend when convrot is also selected. Windows execution requires Windows 10 or 11, a supported NVIDIA GPU with a current driver, and the Visual C++ Redistributable for Visual Studio 2015-2022; a separate CUDA toolkit or Visual Studio install is not required.

What's in here

Why use this

You have multiple PyTorch models that don't all fit on GPU simultaneously, and you want to swap them in and out efficiently across many calls. Re-loading from disk every call is too slow (seconds per gigabyte). Keeping all models resident on GPU is too expensive. torch.cuda.empty_cache() plus .to("meta") gets you the basics but leaves significant performance on the table — pinned host memory does CPU↔GPU DMA at full PCIe bandwidth (~30 GB/s vs. ~3 GB/s from disk), and a shared cache lets multiple models use the same host-memory budget.

This library gives you:

1. Cached resources that pin reusable model or LoRA state to host RAM.
2. Activation lifecycles that move one cached model onto a compute device.
3. A resource cache that evicts least-recently-used unleased entries and protects leased stores from eviction.
4. A model-aware cache that leases model and LoRA resources and owns their device lifecycle.

When to use what

Quick start: cached model use

import torch
from piper_offload import ModelCache, ModelSpec

cache = ModelCache(max_cache_bytes=24 * 1024**3)
model_spec = ModelSpec(
    key="main",
    estimated_cache_bytes=12 * 1024**3,
    factory=build_my_model,  # returns a fresh nn.Module
)
device = torch.device("cuda")

# First use builds and leases the runtime.
with cache.use(model_spec, device=device) as gpu_model:
    output = gpu_model(input_tensor)

with cache.use(model_spec, device=device) as gpu_model:
    output = gpu_model(input_tensor_2)

ModelCache inherits the complete ResourceCache API, adding model activation and LoRA coordination to the same registry and memory budget. ModelSpec factories should build fresh modules. One model cache entry contains one ModelOffloader and one model instance. Uses are sequential: an overlapping activation raises ModelRuntimeInUseError. Applications that need concurrent replicas must register separately constructed models under distinct cache keys, which intentionally duplicates their pinned host storage. To release pinned host memory, evict or clear inactive cache entries and drop any escaped model references.

Pinned copies versus adopted model backing

Model weights use owned pinned host copies by default for maximum CPU-to-GPU transfer bandwidth. Systems with a tight pinned-memory limit can instead adopt existing CPU backing without copying it:

model_spec = ModelSpec(
    key="main-adopted",
    estimated_cache_bytes=12 * 1024**3,
    factory=build_my_model,
    block_paths=("transformer_blocks",),
    host_backing="adopt",
)

# The same option exists on the low-level API.
offload = ModelOffloader.from_module(
    model,
    block_paths=("transformer_blocks",),
    host_backing="adopt",
)

Adopt mode is a strict zero-copy path for frozen model state that already lives on CPU. The offloader retains the existing tensors rather than cloning them, so anonymous pageable allocations and file-backed/mmap storage use the same implementation and preserve their original backing. Copies go directly into the existing GPU target pool; CUDA owns any implicit staging. Trainable parameters, non-CPU tensors, incompatible layouts, and adapters that cannot preserve the source representation raise instead of silently materializing a copy. The policy applies to both streamed blocks and the non-streamed model remainder. LoRA resources choose their host storage independently (pinned by default), and cache_bytes accounts for model bytes the same way in either mode using each tensor adapter's logical byte count. For mmap backing, this is the size of the retained model tensors, not the checkpoint file size or the mmap pages currently resident in RAM.

Because adopted model backing aliases its source allocations, callers must not mutate retained parameter or buffer references—or writable mmap contents— for the offloader's lifetime. requires_grad=False disables autograd tracking; it does not make the underlying bytes immutable.

Adopted capture retains every selected tensor before changing any module registry. Only after the complete composite store succeeds does binding install the adopted wrappers. An adoption failure therefore leaves the supplied model's original parameter and buffer objects untouched.

Pinned backing remains faster for transfer-bound streaming, while adopted backing has near-zero construction-copy cost and can be competitive when computation hides weight movement. Compare the complete offloader path on a target system with:

python benchmarks/benchmark_offloader_host_backing.py

Manual offloader lifecycle

Use ModelOffloader directly when you want explicit lifecycle control without ModelCache.

import torch
from piper_offload import ModelOffloader

model = build_my_model()
offload = ModelOffloader.from_module(model)
device = torch.device("cuda")

offload.activate(device)
try:
    output = offload.value(input_tensor)
finally:
    offload.deactivate()

del offload, model  # drop refs to free pinned host memory

ModelOffloader.from_module() mutates the source model while pinning: frozen nn.Parameter registry entries get repointed at Parameters wrapping pinned CPU storage, trainable Parameter objects keep their identity and point their .data at pinned CPU storage, and buffers are replaced with pinned copies. After construction, only access the bound model while the offloader is active: call activate(device), use offload.value, and guarantee a matching deactivate() with try/finally. For CUDA training, wrap optimizer.step() in offload.optimizer_step() so trainable GPU updates are copied back to the pinned CPU cache before deactivation. Drop the offloader and model references to release pinned host memory — there's no close(); resource cleanup is reference-drop + GC.

Manual block streaming

For models too big to fit on GPU even when active. Streams transformer blocks through a small GPU-resident window using forward-pre hooks and a CUDA-stream-based async prefetcher.

import torch
from piper_offload import ModelOffloader

# Construction pins and binds once; cache_bytes is final immediately.
# block_paths selects what streams.
offload = ModelOffloader.from_module(
    model,
    block_paths=["transformer_blocks"],  # path(s) to the nn.ModuleList
)
device = torch.device("cuda")

offload.activate(device)
try:
    output = offload.value(input_tensor)
finally:
    offload.deactivate()

del offload, model  # drop refs to free pinned host memory

Ordinary streaming owns one active block and one asynchronous lookahead target. That fixed two-target window overlaps the next whole-block copy without retaining blocks that a sequential traversal will reload anyway. Direction changes and iteration wraparound are detected internally. Streaming itself is selected by block_paths and transient_block_paths; with both lists empty, the whole model is one bulk-pinned component that activation copies to the GPU. For heterogeneous block lists, execution still limits concurrency to the active and lookahead blocks, while the morphing pool may park one reusable target per distinct tensor-layout signature.

ModelOffloader only streams on CUDA. Activating the binding on cpu is a pass-through over the already-installed pinned CPU storage: no target pool, no streaming hooks, no weight copies. lora_mode="merge" is CUDA-only; use routed LoRA mode for CPU activation. Routed LoRA installs target-Linear hooks: a forward-PRE hook copies that target's host-backed factors to the input device, and a forward-POST hook applies the residual and releases those device copies.

Optional transient residency

Inference workloads can release large modules and streamed pools as soon as their traversals finish:

offload = ModelOffloader.from_module(
    model,
    transient_block_paths=["transformer_blocks"],
    transient_paths=["input_embedder", "output_head"],
)

Each transient_paths module recursively owns a separate CUDA working set for its non-streamed state. Its successful forward releases that set immediately, before later model work. Every entry in transient_block_paths is a streamed block group whose pool releases after its final block. Ordinary block_paths pools remain resident for the activation. A successful root-model forward reacquires all released components for the next invocation. A conditionally skipped component remains acquired.

The release boundaries are deliberately explicit. A transient_block_paths group releases after its final block and therefore cannot traverse that stack again later in the same root call. For transient_paths, the named module's own forward is the boundary, so its state cannot be used again afterward. A skipped path remains acquired. ModuleList/ModuleDict containers are not expanded for transient_paths, and ModelOffloader does not inspect repeated calls, functional parameter access, aliases across component boundaries, or autograd state. Paths must own disjoint state, and shared-storage aliases must not cross component boundaries; these guarantees belong to the caller. CPU activation remains eager and installs no scheduling hooks. Ordinary and rolling transient block groups both stop at the final block instead of filling block 0 immediately before release.

Every module object in a transient_block_paths group must be distinct. An ordinary block_paths group may still reuse a module object, but a module hook cannot distinguish which aliased list position just completed and therefore cannot provide a safe early-release boundary.

Optional streamed-block compilation

Repeated streamed blocks can opt into forward-only Inductor compilation at runtime construction:

from piper_offload import BlockCompileConfig, ModelOffloader
from piper_kernels.linear.convrot import convrot_int8_compile_options

offload = ModelOffloader.from_module(
    model,
    block_paths=["transformer_blocks"],
    block_compile=BlockCompileConfig(
        dynamic=True,
        fullgraph=False,
        options=convrot_int8_compile_options({"max_autotune": True}),
    ),
)

block_compile=None (the default) preserves eager behavior. One configuration applies to every block_paths and transient_block_paths group and has no effect when neither is configured. The backend remains fixed to Inductor. Optional backend settings can be supplied through options; the mapping is copied for each block before it is forwarded to torch.compile. This is also the boundary for compiler extensions such as Piper Kernels' ConvRot preparation-sharing and activation-folding pass shown above. When no options are provided, compilation uses Inductor's default mode. Piper Offload does not pass mode because PyTorch treats mode and options as mutually exclusive.

Only each distinct block module's forward is compiled. Its module __call__ and Piper Offload's forward-pre hook stay eager, so block activation and prefetch finish before compiled computation begins. Compiled forwards are installed only for CUDA activations and the exact original forwards are restored on deactivate or activation rollback. CPU activation remains eager. The lazy compiled callables are retained by the bound runtime, so later eligible activations can reuse their compiled graphs.

Compilation is inference-only in this initial implementation. Streamed training remains available without block_compile, but combining block compilation with autograd/checkpointed training is unsupported until it has dedicated correctness coverage.

Merge-mode LoRA remains compatible. If any routed LoRA is active, every streamed block runs eagerly for that activation because routed child-Linear hooks stage parameters inside the block forward. The bypass is temporary: a later activation with no routed LoRA uses compiled forwards again. Selecting lora_mode="routed" without supplying a LoRA does not bypass compilation.

An experimental rolling mode can replace the ordinary resident/prefetch block targets with one shared parameter target:

offload = ModelOffloader.from_module(
    model,
    block_paths=["transformer_blocks"],
    block_compile=BlockCompileConfig(rolling=True, fullgraph=True),
)

offload.activate("cuda")
try:
    output = model(inputs)
finally:
    offload.deactivate()

After a parameter's final compiled-graph reader launches, a CUDA event lets a private stream refill that same storage from the next block while the remainder of the current block computes. Immediately before that parameter's first reader in the next block, the compiled graph waits only for its corresponding refill; it does not stall for every parameter in the block. The rollover pass is appended after user/Piper Kernels post-grad graph passes, so its first/last-use analysis sees their rewritten graph. Waits and refills are non-mutating ordered host effects with late scheduler-only ordering edges. They neither consume reader tensors nor model an immutable parameter as mutated, avoiding forced intermediate materialization while preserving the ordinary fusion, memory coalescing, and kernel-autotuning plan. For adapters that support merge-mode LoRA, merge hooks run after every base refill on the copy stream. GGUF and TorchAO INT4 tile-packed weights retain their existing no-merge restriction.

Rolling deliberately fails closed outside its tested contract: fullgraph=True, frozen regular dense, TorchAO-family, Quanto, GGUF, or Piper ConvRot INT8 parameters, homogeneous block layouts, distinct block modules, no tied streamed parameters, and no streamed buffers. Bitsandbytes, DTensor, unreviewed external adapters, and heterogeneous block layouts continue to use the existing morphing block-target pool. Structured logical weights are tracked across every AOT-flattened storage input, and the refill is placed after the last reader of any storage tensor. Repeated resident traversal rolls the final block directly into block zero. Transient streaming stops at the final block and reacquires a fresh block-0 target after the root model forward. Skipped or out-of-order traversal remains correct through a foreground-refill fallback.

For the supported contract, the scheduler-only lifecycle edges preserve the ordinary compiled compute kernels and their autotuning identity. The benchmark checks output equivalence in addition to latency and memory; model-level validation remains prudent after changing PyTorch, compiler extensions, or custom kernels.

Use python benchmarks/benchmark_rolling_compile.py to compare steady-state latency and CUDA allocator residency against ordinary compiled block prefetch on the current GPU. The benchmark also checks output equivalence.

Compiler and backend failures propagate normally. Piper Offload never catches a failed compiled invocation and retries that same block eagerly: with graph breaks, earlier segments may already have executed, so retrying could duplicate mutations. Dynamo's native graph-break and recompilation-limit behavior still applies.

By default, trainable parameters (e.g. LoRA adapters) are managed by the composed PinnedComponent: they move to GPU on CUDA activation and back to pinned CPU storage on deactivate. On CPU activation they stay in the host-backed module state. Wrap CUDA optimizer updates in offload.optimizer_step() so updated trainable bytes are copied back to pinned CPU storage before deactivation.

To reduce trainable-weight residency during training, opt into streaming in-block trainable weights:

offload = ModelOffloader.from_module(
    model,
    block_paths=["transformer_blocks"],
    stream_trainable_weights=True,
)

During CUDA activation in this mode, only the trainable parameter .data streams. It is GPU-resident while its block is resident, plus during the optimizer update. CPU activation remains pass-through. Gradients are not streamed; PyTorch owns param.grad normally.

LoRA merge

ModelOffloader supports optional per-weight LoRA merging through activation arguments. Merge mode installs activation-scoped post-copy hooks for managed parameter targets. Unknown targets raise during activation by default. A LoRA that is intentionally shared across separately loaded model components can set allow_partial_targets=True; its merge and routed uses then apply only the intersection of adapter targets and model parameters, including a valid no-op when that intersection is empty. Present targets still receive the ordinary shape and capability validation. LoRA target keys must match the model's parameter names exactly; any remapping — stripping a diffusion_model. prefix, inserting a PEFT .base_layer. segment — is the caller's job when building the LoRA state dict. Each hook runs immediately after the owning component copies the base weight from pinned CPU storage to GPU, so both block-streamed and non-block weights use the same merge path. Merge compatibility is adapter-owned: plain dense tensors opt into in-place addmm_; structured quantized wrappers can opt into an adapter-owned staged merge, whose implementation may select a format-specific kernel or its own framework-operator fallback. Use routed mode for formats that do not expose either merge capability but still provide a compatible logical Linear weight shape and compute dtype. PinnedParam remains a storage primitive; LoRA merge mode asks the selected adapter for the required update capability.

The LoRA request is scoped to one activate() call. Target lookup is resolved during activation; target compatibility can be preflighted with LoRATransform.validate_target() or validated when the merge hook applies.

For backward compatibility, older PEFT/Wan adapters may include an optional <module>.lora_B.bias vector alongside their A/B matrices. LoRA detects and stores that vector automatically; modern A/B-only adapters keep the existing path unchanged. Merge mode applies the vector to the module's existing plain dense bias with the same adapter strength. It raises before mutation when the legacy vector targets a bias-less module instead of materializing a new model parameter. Direct LoRATransform use treats this as one logical operation via validate_target(weight, bias) followed by apply(weight, bias); validation and application remain explicit phases. Offloaded activation uses separate weight and bias post-copy callbacks so neither update can be overwritten by a later parameter copy. Routed mode can apply the vector as part of its output residual even when the base nn.Linear has no bias. Native Wan/ComfyUI .diff_b conversion remains the caller's responsibility.

import torch
from piper_offload import ModelOffloader, LoRA
from safetensors.torch import load_file

offload = ModelOffloader.from_module(
    model,
    block_paths=["transformer_blocks"],
    # Default: stream_trainable_weights=False
)
device = torch.device("cuda")

# Each LoRA owns immutable factors shared by merge and routed uses.
lora_a = LoRA.from_state_dict(
    state_dict=load_file("lora_a.safetensors"),
)
lora_b = LoRA.from_state_dict(
    state_dict=load_file("lora_b.safetensors"),
)

offload.activate(
    device,
    loras=[lora_a, lora_b],
    lora_strengths=[0.8, 0.5],
    lora_mode="merge",
)
try:
    output = offload.value(input_tensor)
finally:
    offload.deactivate()

LoRA entries are ordered contributions. Repeating the same LoRA applies it again, and each occurrence uses the corresponding lora_strengths value.

Quantized merge uses stochastic rounding by default so LoRA updates smaller than one quantization step are not systematically rounded away. Opt into deterministic round-to-nearest when exact deterministic codes are required:

offload.activate(
    device,
    loras=[lora_a],
    lora_strengths=[0.8],
    lora_mode="merge",
    stochastic_rounding=False,
)

Sampling is an internal LoRA-merge detail. A scalar seed is derived from the full parameter path and that transform's merge count, then used by backend-local randomness without consuming PyTorch's global RNG. Reapplying a streamed merge therefore uses a fresh deterministic sample each time. DTensor additionally derives a seed from each shard's global offsets, while replicated ranks retain matching samples. All LoRAs for a target are accumulated and rounded once. Dense targets still use their exact addmm_ update. Routed mode ignores the option because it never requantizes the base.

derive_seed(*parts) is the public canonical derivation utility used by this path. It accepts strings and unsigned 64-bit integers and is useful when an external adapter needs a deterministic sub-seed:

from piper_offload import derive_seed

local_seed = derive_seed(parent_seed, shard_offset)

LoRA factors use pinned storage by default. To retain existing anonymous pageable or mmap/file-backed CPU tensors without copying them, use strict adoption:

lora = LoRA.from_state_dict(
    state_dict=load_file("lora.safetensors"),
    host_backing="adopt",
)

The loader must return tensors that still reference the desired backing. Adopted LoRA backing requires contiguous CPU factors and raises rather than moving or normalizing them. A dtype= conversion is also rejected because it would allocate new storage; pre-convert the checkpoint or use pinned backing. Because the resource aliases adopted tensors, callers must not mutate those storages during the LoRA's lifetime. Pinned factors generally transfer faster, especially in routed mode where the matched factors move on every invocation.

Block reload from pristine pinned CPU storage automatically clears the previous merge — no explicit unmerge step needed.

Pass lora_mode="routed" as an alternative to the default merge mode. Routed mode installs a forward hook pair on each matched nn.Linear parent — y = base(x) + alpha * (B * A * x + bias) when a legacy bias is present — instead of merging into the base weight. Its PRE hook copies only that target's adapter tensors from pinned CPU storage to the invocation's input device; its POST hook applies the residual and releases those device tensors. Multiple LoRAs on one target are grouped into one hook pair and summed independently. Routed mode is inference-only: factors are frozen (requires_grad=False) and no gradient flows to them. LoRA backing is immutable, so merge and routed uses may overlap across model runtimes. Use routed mode when:

• The base weight is quantized or otherwise structured, but still exposes a logical nn.Linear weight shape and compute dtype, and its adapter does not support merge updates. lora_mode="routed" works because it doesn't touch the base.
• You want to switch LoRAs frequently without re-streaming the underlying base weight or retaining the whole adapter on GPU.

Routed mode is restricted to nn.Linear parents. It handles tied weights by hooking only the exact parent module named by the target, so it never mutates shared storage. Packed formats whose parameter shape differs from the logical matmul weight need a per-format route layer.

For a one-shot permanent merge — bake the LoRA into the model weights and discard the LoRA — use merge_lora:

from piper_offload import merge_lora, LoRA

merge_lora(
    model,
    [(LoRA.from_state_dict(state_dict=load_file("lora.safetensors")), 0.8)],
    stochastic_rounding=False,  # optional deterministic opt-out
)

This uses an in-place addmm_ for plain fp/bf bases. Supported quantized adapters use format-specific Triton kernels on CUDA and retain their dequantize/requantize reference path as a fallback. Formats without a merge path (GGUF, TorchAO INT4 tile-packed) need routed LoRA instead. See Quantized weight support for the full matrix. Unlike an activation-scoped LoRA request, this is not reversible. Unknown targets raise, and all target names, factor shapes, and advertised merge capabilities are preflighted before mutation. Multiple LoRAs for one quantized parameter are packed into one staged low-rank update and the weight is re-encoded once.

For the default stochastic merge, each adapter first uses its existing upstream recipe to compute the final data-dependent scales and other quantization parameters, then samples only the terminal weight code between the two neighboring values on that finalized grid. Exact endpoints and saturation retain the upstream code. Exact-zero strengths are discarded before target lookup or factor staging. Standard CUDA layouts use the same format-specific Triton merge kernels for deterministic and stochastic rounding. Random samples are keyed by logical element index, so launch geometry does not change the result. The Torch and Triton backends replay independently for a fixed seed but do not promise byte-identical samples across implementations or Triton versions. Nested bitsandbytes 4-bit scales still use the reference path because their final effective scale is known only after double quantization. Piper ConvRot INT8 forwards the derived seed to piper-kernels' public addmm_.

This is one composable requantization pipeline per format rather than parallel deterministic and stochastic implementations. Each concrete adapter's existing requantize(t, like=..., rounding_seed=None) first constructs its normal finalized representation; when rounding_seed is supplied, it then recodes only terminal weight data against those stored parameters. Omitting the keyword preserves the original deterministic bytes. The public structural conversion protocol retains its deterministic minimum signature for downstream static compatibility.

Heterogeneous block lists

block_paths and transient_block_paths accept dotted paths for models with multiple kinds of blocks (e.g. Flux's transformer_blocks + single_transformer_blocks). Each path becomes its own streaming group with its own target pool; the two lists are mutually exclusive. Blocks within a group must share parameter and buffer names plus trainability structure, but their shapes, dtypes, quantization formats, alias topology, and buffer layouts may differ. The morphing pool keys reusable targets by those layouts. For bespoke grouping, compose StreamedComponentStore instances directly:

offload = ModelOffloader.from_module(
    model,
    block_paths=["transformer_blocks", "single_transformer_blocks"],
)

Training streamed blocks

Training through a streamed block requires activation checkpointing on each block — wrap call sites in torch.utils.checkpoint.checkpoint, or call model.gradient_checkpointing_enable() on a HuggingFace model. Without it, loss.backward() raises:

RuntimeError: one of the variables needed for gradient computation
has been modified by an inplace operation

The reason is autograd's saved-tensor mechanism. A Linear saves a reference to its weight tensor at forward time and records the tensor's version counter. Streaming is a sequence of in-place copy_ writes into a fixed pool of GPU target tensors — every block load bumps the target tensor's version, so by the time backward arrives at an earlier block, the target has been overwritten and the version mismatch raises.

Activation checkpointing sidesteps this. With checkpointing, the block's internal forward runs under no_grad — no internal tensors are saved for backward. When backward arrives, PyTorch re-runs the block's forward with grad enabled, building a fresh autograd graph whose saved references only live within that one block's recompute-then-backward window. Target reuse outside that window is safe because no autograd graph spans across reuses.

import torch
from piper_offload import ModelOffloader

offload = ModelOffloader.from_module(
    model,
    block_paths=["transformer_blocks"],
)
device = torch.device("cuda")

model.gradient_checkpointing_enable()  # required for training
model.train()

offload.activate(device)
try:
    gpu_model = offload.value
    for batch in loader:
        loss = gpu_model(**batch).loss
        loss.backward()
        with offload.optimizer_step():
            optimizer.step()
        optimizer.zero_grad()
finally:
    offload.deactivate()

Checkpointing every streamed training block is the caller's responsibility — ModelOffloader does not auto-detect or warn about its absence. It matters most with stream_trainable_weights=True, where the .data swap bypasses autograd's version-counter check, so missing checkpointing can silently corrupt gradients. Verify every streamed training block is checkpointed (HF gradient_checkpointing_enable() or manual torch.utils.checkpoint.checkpoint wrapping).

Wrap CUDA optimizer updates so managed trainable weights are synced back to pinned CPU storage. With stream_trainable_weights=True, this also materializes streamed trainable weights on GPU while a normal PyTorch optimizer mutates them:

offload.activate(device)
try:
    gpu_model = offload.value
    for batch in loader:
        loss = gpu_model(**batch).loss
        loss.backward()

        with offload.optimizer_step():
            optimizer.step()

        optimizer.zero_grad()
finally:
    offload.deactivate()

This boundary is not optimizer-specific. It runs whatever optimizer.step() does, copies updated trainable data back to pinned CPU storage, and leaves gradients on GPU.

Cached model details

ResourceCache owns only reusable-resource admission, accounting, leases, and eviction. ModelCache inherits that API and adds dependency leasing, LoRA attachment, and device activation for model uses.

from piper_offload import (
    LoRASpec,
    ModelCache,
    ModelSpec,
)
from safetensors.torch import load_file

cache = ModelCache(max_cache_bytes=80 * 1024**3)
device = "cuda:0"

text_encoder = ModelSpec(
    key="text_encoder",
    estimated_cache_bytes=12 * 1024**3,
    factory=build_text_encoder,
)
diffusion_model = ModelSpec(
    key="diffusion_model",
    estimated_cache_bytes=24 * 1024**3,
    factory=build_diffusion_model,
    block_paths=("transformer_blocks",),
)
style_lora = LoRASpec(
    key="style-lora",
    estimated_cache_bytes=512 * 1024**2,
    factory=lambda: load_file("style.safetensors"),
    dtype=torch.bfloat16,
)

with cache.use(text_encoder, device=device) as enc:
    embeddings = enc.encode(prompt)

with cache.use(
    diffusion_model,
    device=device,
    lora_specs=[style_lora],
    lora_strengths=[0.8],
    lora_mode="routed",
) as model:
    latent = model(...)

The model cache leases LoRA resources before admitting the model resource. An adapter selected for a use therefore cannot be evicted by that same model admission. All leases unwind in reverse order if construction or activation fails. lora_strengths defaults to 1.0 per LoRA; when supplied, it must have the same length as lora_specs. Exact 0.0 and -0.0 strengths are inactive: they are filtered before target grouping or hook installation, and ModelCache does not construct or lease their LoRA resources. Merge and routed uses may share one cached LoRA across model runtimes or repeat it within one use because each occurrence is an independent contribution with its own strength.

For direct resource access, use a cache lease:

with cache.lease(style_lora) as lora:  # auto-registers on first lease
    targets = lora.targets

Anti-pattern: the factory should build a fresh model each call, not capture an externally-held one. With factory=lambda: my_kept_model the cache is no longer the sole owner of the model. Always have the factory build the model itself.

ResourceCache accepts custom EvictionPolicy implementations. The default is LRUEvictionPolicy for unleased host-cache eviction. The cache builds the eviction candidate set and byte context, then asks the eviction policy to choose victims; ResourceCache still owns validation, accounting, admission, and release. Policies are called under the cache lock. choose_victims() must return unique keys from context.candidates and enough bytes to satisfy context.bytes_to_free; otherwise ResourceCache raises EvictionPolicyError without evicting anything.

Architecture

registration / cache admission
------------------------------
            ResourceSpec protocol
                    |
         ModelSpec / LoRASpec / ObjectSpec
        |
        v
  +-------------+
  | ModelCache  |  extends ResourceCache with model-aware use
  +-------------+
        |
        +-- builds/admit --> ModelOffloader (one model, one runtime)
        |                    |
        |                    +-- PinnedComponent
        |                    |       |  resident non-block state
        |                    |       +-- PinnedParam(s)
        |                    |
        |                    +-- PinnedComponent(s)
        |                    |       |  transient path state
        |                    |       +-- PinnedParam(s)
        |                    |
        |                    +-- StreamedComponent(s)
        |                            |
        |                            +-- PinnedParam(s)
        |
        +-- builds/admit --> LoRA (pinned or adopted factors)
        |
        +-- builds/admit --> custom ResourceStore
        |
        +-- chooses inactive victims via EvictionPolicy

ModelCache.use(...)
-------------------
ModelCache
   |
   +-- lease LoRASpec(s), then ModelSpec
   +-- ModelOffloader.activate(loras=...) claims the model runtime
   +-- yield ModelOffloader.value
   +-- ModelOffloader.deactivate() removes LoRA hooks + releases model

ResourceSpec is the structural registration contract: key, estimated_cache_bytes, build_store(), and value(store). The standard specs are independent frozen dataclasses; custom specs can implement the protocol without inheriting from them. ResourceStore is the backing-state contract and reports cache_bytes. A cache lease protects that store from eviction but does not create or activate a runtime. ResourceBinding is the active-resource lifecycle contract: value, activate(device=None, **kwargs), and deactivate(). ModelOffloader is both a cached ResourceStore and a ResourceBinding; LoRA is an immutable cached ResourceStore. It exposes neither an active lifecycle nor a model-like value; merge and routed hooks read its pinned factor backing directly.

A custom cached resource needs only one spec and one store:

from dataclasses import dataclass

from piper_offload import ResourceSpec, ResourceStore


class MyStore:
    @property
    def cache_bytes(self) -> int: ...


@dataclass(frozen=True)
class MySpec:
    key: str
    estimated_cache_bytes: int

    def build_store(self) -> MyStore:
        return MyStore()

    def value(self, store: ResourceStore) -> MyStore:
        assert isinstance(store, MyStore)
        return store


spec: ResourceSpec[MyStore] = MySpec(
    key="my-resource", estimated_cache_bytes=...,
)

with cache.lease(spec) as store:
    ...

StreamedComponent and PinnedComponent are composable activate/deactivate lifecycle pieces (no value or model) that live inside a top-level model runtime rather than acting as one themselves. Either active component may release() and later acquire() its CUDA working set without ending the activation session; activation acquires it immediately by default. ModelOffloader.register_forward_hook() registers a native PyTorch forward hook by fully-qualified module name and returns a caller-owned remover. This lets higher-level runtimes coordinate component lifetimes at model execution boundaries without adding policy to component internals.

TensorAdapter is the per-parameter extension point. Its base contract only covers inference movement: clone/pin, H2D copy, GPU wrapper rebuild, cache bytes, logical compute dtype, and block-layout signatures. Extra behaviors are explicit capabilities: CPU round-trip for optimizer-step sync, Parameter.data swap for trainable streaming, shape-preserving dequantize/requantize conversion, representation-preserving copy_into, and adapter-owned staged LoRA merge. Zero-copy host adoption is another optional capability: adopt_host() returns adapter-specific state that aliases the existing CPU storage. The existing cache_bytes() method accounts for either pinned or adopted state. LoRA dispatch uses either dense in-place addmm_ for plain bases or the staged merge capability for structured bases; conversion and copy capabilities do not implicitly advertise merge support. Adapters whose merge supports only certain layouts or staged factor values can also implement LoRAMergeValidationTensorAdapter.validate_lora_merge. Permanent merge stages and validates every requested operation through this hook before mutating any weight; DTensor delegates validation to its local-shard adapter. The merge and validation protocols include an optional rounding_seed: int | None = None keyword. Downstream adapters should accept that keyword even when they only implement deterministic rounding; omitting it or passing None preserves deterministic behavior. An adapter that needs a reproducible substream can derive one with the public derive_seed() utility.

Downstream tensor subclasses can provide their adapter without adding a format-specific dependency to piper-offload:

from piper_offload import (
    AdoptableTensorAdapter,
    TensorAdapter,
    register_adapter,
)


class MyTensorAdapter:
    # Implement the stateless TensorAdapter protocol.
    # To support host_backing="adopt", also implement the
    # AdoptableTensorAdapter protocol and return adopted adapter state.
    ...


remove_adapter = register_adapter(MyTensorAdapter)

Register adapters during application startup, before constructing models or pinned resources. DTensorAdapter remains the outermost wrapper; registered adapters are then checked newest-first before the remaining built-ins. This lets a downstream adapter override a built-in isinstance match for a more specific subclass, and also lets DTensor delegate a custom local shard through the same registry. register_adapter() returns an idempotent removal callable for tests and scoped integrations.

Cached resource lifecycle

Cached resources own cache accounting. Host capture happens during construction so cache_bytes is final at admission time; leases protect resources while they are used. ModelOffloader owns one exclusive activation lifecycle. LoRA remains immutable host backing throughout its lease:

ModelOffloader: construct -> lease -> activate <-> deactivate -> release lease
LoRA:            construct -> lease -> read host factors -> release lease

ModelOffloader.activate(device=...) makes the model usable for compute on the requested device. Merge hooks copy factors when their base weight is loaded; routed PRE hooks copy factors for one Linear invocation and routed POST hooks release them after enqueueing the residual. ModelOffloader, MpsWeights, PinnedComponent, and StreamedComponent require an explicit device. CUDA activation uses the streaming/DMA path where applicable; CPU activation is pass-through over pinned host-backed storage. deactivate() releases transient device resources. Host backing remains cached until its resource is evicted or otherwise released.

Construction optimizes peak host memory. Pinning clones managed tensors into pinned CPU storage. For plain torch.Tensor parameters, the source Parameter.data may be immediately repointed at the pinned clone as soon as that pinned parameter is created. This releases the original source storage early, avoiding temporarily holding both pageable and pinned copies for CPU-origin models and promptly freeing GPU storage for CUDA-origin models. It is a clone-to-pinned plus storage swap, not true in-place pinning. Tensor subclasses such as quanto, GGUF, and NVFP4 do not use this .data swap when it would lose wrapper state.

There is no close(). To release pinned host memory, first let all leases end, then evict or clear the cache entry. Python's refcount-based GC frees pinned tensors once the cache and any escaped resource, binding, or model references are gone.

Failure semantics. If construction raises after pinning has started, the model may already be partially repointed to pinned storage. Treat the partially constructed resource/model as unrecoverable: drop those references and rebuild from a fresh model instance. If activate() raises, the offloader rolls back its active components and releases its activation claim, so a later well-formed activation may retry the same cached resource. Routed LoRA hook-registration failures remove any hooks already installed; permanent merge validates all targets before mutation. This is a low-level library; we don't guard against caller misuse.

Compatibility

• torch.compile support is deliberately narrow. Use block_compile=BlockCompileConfig(...) to compile only declared streamed block forwards during CUDA inference. External whole-model torch.compile(model), model.compile(), compilation of the pinned remainder, and compiled streamed training remain unsupported. Routed LoRA temporarily bypasses compiled blocks. Compiler code/artifact caches and compiler-owned workspace are outside ResourceCache.cache_bytes; model eviction does not call process-global torch.compiler.reset(). Experimental rolling=True compilation has the additional homogeneous/full-graph and adapter restrictions documented above.
• Wrap before DDP/FSDP, not after. Those wrappers manage parameter storage themselves and conflict with the registry-replacement pattern.
• One runtime per cached model. ResourceCache serializes resource construction and lease accounting, then releases its lock while a lease is held. Each cached ModelOffloader owns one model and rejects overlapping activation, including calls from different runners. Concurrent replicas require distinct ModelSpec keys and therefore distinct pinned storage.
• Buffer mutations during CUDA activation are discarded on deactivate(). CPU activation is pass-through over host-backed buffers, so CPU buffer mutations behave like ordinary module mutations. Suitable for inference of stateless modules; not suitable for models that need persistent buffer state across calls (BatchNorm running stats updated in training mode, RNN/SSM hidden state, KV cache).
• Training requires activation checkpointing on every streamed block (model.gradient_checkpointing_enable() for HF models, or manual torch.utils.checkpoint.checkpoint wrapping). Without it, loss.backward() raises an in-place modification error from autograd's saved-tensor check. See Training streamed blocks.

Tied weights

ModelOffloader handles the standard tie_weights() pattern (one Parameter referenced under multiple names) plus the rarer case of distinct quanto wrappers around shared inner _data storage.

ModelOffloader is intended for ordinary transformer block lists where the streamed block weights are independent. Shared storage is preserved within one streamed block and within pinned state, including the standard tied input-embedding/output-head pattern. It is unsupported across offload ownership boundaries: between streamed and pinned state, or between distinct streamed blocks or block groups. ModelOffloader does not prevalidate these unusual layouts; omit block_paths and use whole-model offloading if that sharing must be preserved.

Quantized weight support

Every supported weight type can be offloaded (pinned host ↔ GPU movement). LoRA differs by type: a base can be merged into when its adapter exposes an in-place update path, and routed LoRA (lora_mode="routed" — a forward hook) is available for any of them whose owning module is a logical nn.Linear with compatible shape and dtype, no merge capability required.

Notes:

• Stochastic rounding is supported by every merge-capable built-in quantized adapter in the table. Standard CUDA layouts use fused Triton terminal-code selection; unsupported layouts and nested bitsandbytes 4-bit scales retain the dequantize/requantize reference path. Both preserve the same scale, calibration, packing, and wrapper metadata contract. Plain floating-point and DTensor-wrapped dense weights accept the option but have no quantization code to randomize.

• Merging into a quantized base is lossy because the updated value is re-encoded onto the quantization grid; choosing merge vs routed is the caller's accuracy/latency tradeoff, and is coarser the fewer bits the format has (e.g. MXFP4 / NVFP4 at 4 bits). Data-dependent weight scales are recomputed from the merged values. Formats that cannot safely encode with a zero scale, including Quanto qint8/qfloat8, floor exact-zero blocks to a small positive value; bitsandbytes int8 retains an exact zero scale for an all-zero row.

• Triton dispatch is layout-conservative. Each adapter checks its raw storage, scale metadata, compute dtype, and device before launching. A valid but nonstandard representation uses the reference merge when that format can re-encode the layout.

• † MX and NVFP4 store weights in a block-structured packed layout, so the standard re-encode (which produces the contiguous layout) cannot fill a transposed (non-contiguous qdata) target's storage; those raise a clear error pointing to routed LoRA. A transposed scaled-FP8 PerGroup weight is likewise unsupported because TorchAO only reconstructs groups along the last axis. PerRow and PerTensor scaled-FP8 transposes remain mergeable. int8 cannot be transposed.

• ‡ DTensor merge supports rank-two weights with ordinary Replicate and contiguous Shard placements. Each rank selects the rows and columns needed by its local weight shard from the plain host-backed LoRA factors before device staging, then delegates the update to that shard's adapter; no collective is required. The full factors remain in host memory. The inner adapter must support LoRA merge. Unsupported local tensor types and placements must use routed LoRA. DTensor factors themselves are not accepted in merge mode.

• CPU round-trip (D2H, for context-free CPU optimizer steps) and trainable Parameter.data swap are separate capabilities: plain tensors have both; quanto and both scaled-FP8 representations add CPU round-trip; the other quantized formats are movement + (where shown) merge only. See the per-format sections below.

Quanto support

Quanto-quantized models (optimum.quanto.WeightQBytesTensor) are handled correctly by both ModelOffloader modes. PinnedParam decomposes the wrapper into its inner _data (int8/fp8) and _scale (fp16/fp32) tensors, pins each, and reconstructs the quanto wrapper around the GPU storage on activation.

Optimized MarlinF8QBytesTensor weights are first canonicalized to the ordinary unpacked WeightQBytesTensor representation for pinned streaming, so streamed execution is correct but does not use Marlin's packed matmul. Direct merges into an existing Marlin weight use the reference dequantize/addmm/requantize path and repack the result into that weight's original physical storage.

A naive param.data.clone() on a quanto tensor silently dequantizes it via the dispatch fallback — the explicit decomposition is required for correctness.

LoRA on qint8 and the common E4M3/E5M2 qfloat8 bases uses a Triton merge on CUDA that recomputes Quanto's absmax weight scale from the merged values. Other layouts, including qfloat8 E4M3FNUZ, use the equivalent dequantize/addmm/requantize path with the same scale policy. Exact-zero scale blocks are repaired to a small positive value so qfloat8 never encodes a 0 / 0 NaN. Both paths update the existing inner data and scale storage; neither attempts native in-place addmm_ on a WeightQBytesTensor. Use lora_mode="routed" when the base must remain untouched or adapters need to switch without reloading it.

Piper ConvRot INT8 support

Piper ConvRot weights (piper_kernels.linear.convrot.ConvRotInt8Tensor) are handled when the convrot optional extra is installed. piper-kernels owns the tensor semantics plus reference and optimized execution backends; Piper Offload owns only the built-in PiperConvRotInt8Adapter. PinnedParam pins the INT8 qdata and float32 per-output scale, preserves group_size and the logical floating dtype, and reconstructs the same wrapper around CUDA storage on activation.

The adapter remains frozen-only: it does not expose CPU round-trip or trainable Parameter.data swap. Merge-mode LoRA delegates the staged low-rank update and optional reproducible stochastic-rounding seed to Piper's public in-place ConvRotInt8Tensor.addmm_, preserving the wrapper and its storage identities. Piper uses its optimized Triton backend on supported CUDA devices and its portable reference backend elsewhere. Use routed LoRA when the base must remain untouched. The base package remains importable without piper-kernels; use uv sync --extra convrot --group dev and then pytest tests/test_piper_convrot_int8_adapter.py -q -rs to exercise the optional suite.

TorchAO NVFP4 support

TorchAO NVFP4 weights (torchao.prototype.mx_formats.nvfp4_tensor.NVFP4Tensor) are handled when the torchao optional extra is installed. PinnedParam pins the packed FP4 qdata, FP8 block scale, optional per-tensor scales, and the TorchAO dispatch metadata, then rebuilds the same concrete NVFP4Tensor subclass around GPU storage on activation. Subclass identity also survives requantization and merge-mode LoRA. The optional extra requires the package's supported TorchAO release; dynamic NVFP4 matmul execution still depends on Blackwell-class CUDA hardware and the matching PyTorch CUDA stack. For uv-managed installs, this repo routes torch on Linux/Windows and torchao on Linux through PyTorch's CUDA 13.0 wheel index. Windows uses TorchAO's portable PyPI wheel because the CUDA 13.0 index does not publish a Windows TorchAO wheel. Use uv sync --extra torchao --group dev and then pytest tests/test_nvfp4_adapter.py -q -rs to exercise the optional TorchAO NVFP4 coverage.

Contiguous rank-two NVFP4 weights support a block-local Triton merge for ordinary or swizzled scales. Single-level scaling needs one merge/pack pass; two-level scaling first reduces the merged amax to a new global scale and then recomputes and packs each block. Neither path materializes the dense weight. Unsupported layouts use the existing NVFP4Tensor.to_nvfp4 reference path. Both preserve the wrapper and its dispatch metadata while re-deriving the data-dependent scales. Like any merge into a quantized base, NVFP4's 4-bit re-encoding is lossy.

The adapter does not opt into CPU round-trip or trainable Parameter.data swap: the quant state lives in the wrapper object, not its bytes, so NVFP4 weights stay frozen for streaming/training. Routed LoRA remains the non-destructive alternative when the target module is a logical nn.Linear with compatible shape and compute dtype.

TorchAO MX (MXFP8 / MXFP4) support

TorchAO MX (OCP microscaling) weights (torchao.prototype.mx_formats.mx_tensor.MXTensor, created by quantize_(...) with an MX inference config or directly via MXTensor.to_mx) are handled when the torchao optional extra is installed. A single adapter covers both MXFP8 (float8_e4m3fn / float8_e5m2) and MXFP4 (float4_e2m1fn_x2), since TorchAO models them as the same MXTensor subclass parameterized by elem_dtype. PinnedParam pins the packed qdata, the E8M0 block scale, and the TorchAO dispatch metadata (elem_dtype, block_size, kernel_preference, act_quant_kwargs, is_swizzled_scales), then rebuilds the MXTensor wrapper around GPU storage on activation. MXFP6 and any other MX element dtype are not admitted; such a tensor falls through to a clear "no adapter" error rather than being silently mishandled. MX matmul execution still depends on Blackwell-class CUDA hardware and the matching PyTorch CUDA stack. Use uv sync --extra torchao --group dev and then pytest tests/test_mx_adapter.py -q -rs to exercise the coverage.

Standard blocksize-32 MX weights support a block-local Triton merge for MXFP8 E4M3/E5M2 and packed MXFP4, including regular or swizzled scales and TorchAO's FLOOR, RCEIL, CEIL, and EVEN scale modes when the mode is recorded in act_quant_kwargs. Weight-only wrappers do not retain the mode and use TorchAO's default FLOOR when re-encoded. Each kernel program updates and packs one 32-element block without materializing the dense weight. Unsupported layouts use the existing MXTensor.to_mx reference path. MXFP4's grid makes a permanent merge much coarser than MXFP8. The adapter does not opt into CPU round-trip or trainable Parameter.data swap: like NVFP4, the wrapper's quant state lives in the object, so MX weights stay frozen. Routed LoRA remains the non-destructive alternative.

TorchAO scaled FP8 support

TorchAO scaled-fp8 weights (torchao.quantization.Float8Tensor, created by quantize_(..., Float8WeightOnlyConfig/Float8DynamicActivationFloat8WeightConfig)) are handled when the torchao optional extra is installed. PinnedParam pins the fp8 qdata and fp32 scale tensors plus the TorchAO dispatch metadata (block_size, mm_config, kernel_preference, act_quant_kwargs), then rebuilds the Float8Tensor wrapper around GPU storage on activation. Per-group, per-row, and per-tensor scale granularities are supported; fp8 matmul execution requires SM89+ (Ada/Hopper or newer) CUDA hardware.

Standard scaled-FP8 layouts use format-specific Triton merges on CUDA and recompute the affected scales; unsupported layouts retain the public Float8Tensor.from_hp reference path. The GPU representation is byte-identical to the host one, so CPU round-trip is also available. Trainable Parameter.data swap is not — scaled-FP8 weights stay frozen.

TorchAO's calibrated static-activation representation is handled separately by StaticFloat8Adapter. It targets only torchao.prototype.quantization.float8_static_quant.prototype_float8_tensor.PrototypeFloat8Tensor weights with per-tensor weight and activation quantization, and pins the FP8 qdata, weight scale, and checkpoint-provided act_quant_scale. All three are included in identity, block-pool layout compatibility, cache accounting, H2D/D2H movement, and wrapper reconstruction; the ordinary Float8Tensor adapter remains the weight-only/dynamic path.

TorchAO 0.17 normally requires the activation scale rank to equal the input rank. Piper Offload's static adapter installs a narrow nn.Linear dispatch shim that flattens ordinary activations before static quantization and reshapes the result afterwards. A checkpoint scalar (or any one-element scale layout) therefore works unchanged for both 2-D and 3-D Linear inputs. LoRA merge uses a format-specific Triton kernel pipeline on CUDA when Triton is available, independently of block_compile. It fuses dequantization, the low-rank GEMM, addition, and tile-level maximum collection before reducing the new per-tensor weight scale and requantizing. CUDA installations without Triton use the same raw storage through ordinary Torch operations; CPU merges retain the generic adapter path. All paths copy only the re-encoded weight bytes and scale into the target; the calibrated activation scale is preserved exactly. Routed LoRA is supported as the non-destructive alternative. Output activation quantization and non-per-tensor Prototype layouts are outside this adapter's contract and are rejected explicitly.

Failure modes

The cache and bindings surface failures as typed exceptions rather than silent corruption.

State Inspection

Use cache.used_cache_bytes and cache.available_cache_bytes for current cache accounting. Use cache.info(key) for per-key state when needed.

The budget can be changed while the cache is running:

cache.resize(40 * 1024**3)
# Equivalently:
cache.max_cache_bytes = 40 * 1024**3

Growing the budget preserves cached entries. Shrinking evicts inactive entries according to the configured eviction policy. If leased entries make the target size impossible, resizing raises ResourceTooLargeError without changing the previous budget or evicting entries.

License

Licensed under the Apache License, Version 2.0.