`core.pipeline_parallel.pinned_activation_buffer`#

Bounded pinned storage for fine-grained activation offloading.

Module Contents#

Classes#

PinnedActivationBuffer

Pack saved activations into one fixed, reusable pinned buffer.

Functions#

`_cuda_runtime`	Return the cuda-python runtime bindings.
`_free_host`	Free a cudaHostAlloc allocation with the allocating device current.

Data#

logger

API#

core.pipeline_parallel.pinned_activation_buffer.logger#: ‘getLogger(…)’

core.pipeline_parallel.pinned_activation_buffer._cuda_runtime()#: Return the cuda-python runtime bindings.

core.pipeline_parallel.pinned_activation_buffer._free_host(address: int, device: int) → None#

Free a cudaHostAlloc allocation with the allocating device current.

Garbage collection may run this on any thread; selecting the device keeps the runtime from creating a context on another GPU.

class core.pipeline_parallel.pinned_activation_buffer.PinnedActivationBuffer(capacity_bytes: int)#

Pack saved activations into one fixed, reusable pinned buffer.

Allocations append within the buffer. Once all backups have been reloaded, the next allocation reuses the buffer after waiting for its H2D readers. Interleaved chunks share one budget; storage is not reclaimed until all outstanding backups are consumed. No individual-tensor free list is needed.

Captured addresses remain reserved for this object’s lifetime: replay does not execute Python allocation/release bookkeeping. Each capture owns a separate range, reusable within that capture after its readers finish. Eager epochs use the remaining tail and cannot overwrite graph backups.

Capacity may be any positive multiple of 256 bytes and is pinned at that size: the storage comes from cudaHostAlloc directly rather than from the PyTorch pinned host allocator, which would round the request up to the next power of two. Storage is allocated at construction and freed once this object and every view of it are gone. Captured graphs hold raw addresses, so keep this object alive while any graph that uses it can replay. Exhaustion raises an error; the buffer never grows.

Initialization

ALIGNMENT#: 256

_allocate_storage(capacity_bytes: int) → torch.Tensor#

Pin exactly capacity_bytes (rounded up to whole pages) with cudaHostAlloc.

This is the same page-locked, mapped allocation the PyTorch pinned host allocator makes, without its power-of-two rounding. torch.frombuffer keeps the backing object alive while any view of the storage exists, so its finalizer frees the pages only after the last view is gone. The finalizer does not run at interpreter exit: cudaFreeHost synchronizes the device, and process exit releases the pages anyway.

allocate( shape: tuple, dtype: torch.dtype, *, stream: torch.cuda.Stream | None = None, ) → torch.Tensor#

Allocate a view whose first write uses the current stream.

Pass the already-active stream to avoid rediscovering it for every tensor. This does not switch streams; the caller still enqueues the copy on it.

release( tensor: torch.Tensor, *, stream: torch.cuda.Stream | None = None, ) → None#: Retire a backup after its final H2D read is enqueued on the active stream.

core.pipeline_parallel.pinned_activation_buffer#