spdl.io.transfer_tensor_h2d

transfer_tensor_h2d(batch: T, /, *, num_caches: int = 4) → T[source]

Transfer PyTorch tensors from CPU to CUDA in a dedicated stream.

Added in version 0.8.0.

Changed in version 0.9.0: Recursive failures drain submitted copies before unwinding. Failed drains retain source/destination owners until completion is proved. Returned GPU destinations are recorded on the selected device’s default consumer stream; callers using other streams retain normal readiness/recording responsibilities.

This function wraps calls to torch.Tensor.pin_memory() and torch.Tensor.to(), and executes them in a dedicated CUDA stream.

When called in a background thread, the data transfer overlaps with the GPU computation happening in the foreground thread (such as training and inference).

See also

Multi-threading (custom) - An intended way to use this function in Pipeline.

benchmark_pinned_arena - Matched public pool/ring pinning comparisons through subprocess iteration and synchronized H2D wall time, with archival earlier measurements.

../_static/data/parallelism_transfer.png

Concretely, it performs the following operations.

  1. If a dedicated CUDA stream local to the calling thread is not found in thread-local storage, creates and caches one. (The target device is determined by the "LOCAL_RANK" environment variable.)

  2. Activates the CUDA stream.

  3. Traverses the given object recursively and transfers CPU tensors to CUDA. pin_memory reuses source storage when Torch recognizes it as pinned; otherwise it creates a page-locked copy. CUDA transfer is submitted with .to(non_blocking=True).

  4. Synchronizes the stream, to ensure that all the data transfers are completed.

Parameters:

  • batch – A torch.Tensor or a composition of tensors with container types such as list, tuple, dict and dataclass.

  • num_caches –

    Number of batch caches to maintain the reference to. This parameter helps mitigate race conditions when using multi-threading with multiple CUDA streams.

    See PyTorch CUDA Race Condition in Multi-threading for details on the rationale behind this parameter.

Returns:

An object of the same type as the input, but the PyTorch tensors are transferred to CUDA device.

Example

from concurrent.futures import ThreadPoolExecutor

from spdl.io import transfer_tensor_h2d

with ThreadPoolExecutor(max_workers=1) as executor:
    future = executor.submit(transfer_tensor_h2d, cpu_batch)
    cuda_batch = future.result()