Skip to content

Layerwise Offload

Layerwise offload streams transformer blocks through one accelerator. For commands and model support, see the layerwise user guide.

Hook ring

LayerWiseOffloadBackend discovers every streamable block and registers one LayerwiseOffloadHook per block. Hooks form a ring: each hook owns the pinned host representation of its registered block and a reference to the next block it must prefetch; the final hook prefetches the first block for the next denoising iteration.

During initialization, parameters and buffers are grouped by dtype and copied into contiguous pinned CPU tensors. Original tensor storage is replaced with an empty or meta placeholder while identity and module structure are retained. Metadata records names, offsets, shapes, and dtypes for rematerialization.

Forward lifecycle

For each block:

  1. the previous hook's completion event ensures this block is materialized;
  2. the copy stream asynchronously allocates and fills the next block's device tensors;
  3. parameter and buffer storage is rebound to views of those tensors;
  4. the current block computes on the normal stream; and
  5. its post-forward hook replaces device storage with placeholders after all dependent work is safe.

The copy stream waits on the compute stream before reusing storage, and the compute stream waits on the per-block prefetch event before forward. These dependencies, not incidental global synchronization, define correctness.

Topology contract

A DiT declares ordered block-container attributes using _layerwise_offload_blocks_attrs or an OffloadPlan. A text encoder declares its block-list paths through OffloadPlan.encoder_block_attrs. Each resolved item must be an executable nn.Module; ordering must match forward execution. Multiple containers are streamed as separate hook rings.

Non-block state of a selected DiT or text encoder remains device resident. Unselected components, VAEs, and declared resident modules are not streamed by this backend.

Invariants and limitations

  • Each offloaded tensor has one pinned host backing store.
  • Parameter identity is preserved when storage is rebound.
  • Blocks with heterogeneous shapes and dtypes use their recorded metadata; implementations must not assume identical blocks.
  • The ring is initialized before the first forward so block zero is available.
  • This backend uses rank-local host weights. Multi-device compute may remain active, but cross-rank weight reconstruction belongs to Distributed Layerwise Offload.

Shared selection and lifecycle behavior is defined in the CPU Offloading design.