nemo_rl.utils.cuda_ipc#

Compatibility shim for CUDA IPC storage handles crossing torch versions.

NeMo-RL hands CUDA tensors from the training venv to the inference venvs (sglang, TRT-LLM, vLLM) through torch.multiprocessing.reductions IPC handles. The storage handle torch produces is a small byte string written by c10::cuda::CUDACachingAllocator::shareIpcHandle::

[SHAREABLE_HANDLE_VERSION][type byte][payload]

where type is 'c' (a plain cudaMalloc block, followed by the 64-byte cudaIpcMemHandle_t) or 'e' (an expandable segment). The consumer rejects any handle whose version byte is newer than its own with received sharable handle from a future version of torch that this version does not know how to handle.

torch 2.13 bumped the version byte from 2 to 3 without changing the 'c' payload, so when the training venv runs torch 2.13 and an inference venv runs torch 2.11 (sglang / TRT-LLM, whose kernels are built against 2.11) every colocated refit dies at the first handle. Rewriting the version byte of 'c' handles to the legacy value makes them readable by both versions. The expandable-segment format did change, so 'e' handles – what the trainer emits under PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True – are left alone; they still open on a same-torch consumer (vLLM) and trip torch’s version check on a torch 2.11 one (sglang / TRT-LLM).

Module Contents#

Functions#

normalize_cuda_ipc_handle

Rewrite a newer 'c'-type storage handle so older torch can open it.

Data#

CUDA_IPC_MEM_HANDLE_SIZE

sizeof(cudaIpcMemHandle_t); a raw handle of exactly this size predates versioning.

SHAREABLE_CUDA_MALLOC

Type byte for a plain cudaMalloc block (SHAREABLE_CUDA_MALLOC in c10).

LEGACY_SHAREABLE_HANDLE_VERSION

The version byte torch <= 2.12 writes. torch 2.13 writes 3 with an identical 'c' payload.

CURRENT_SHAREABLE_HANDLE_VERSION

The version byte torch 2.13 writes; its 'c' payload is identical to version 2’s.

API#

nemo_rl.utils.cuda_ipc.CUDA_IPC_MEM_HANDLE_SIZE#

64

sizeof(cudaIpcMemHandle_t); a raw handle of exactly this size predates versioning.

nemo_rl.utils.cuda_ipc.SHAREABLE_CUDA_MALLOC#

‘ord(…)’

Type byte for a plain cudaMalloc block (SHAREABLE_CUDA_MALLOC in c10).

nemo_rl.utils.cuda_ipc.LEGACY_SHAREABLE_HANDLE_VERSION#

2

The version byte torch <= 2.12 writes. torch 2.13 writes 3 with an identical 'c' payload.

nemo_rl.utils.cuda_ipc.CURRENT_SHAREABLE_HANDLE_VERSION#

3

The version byte torch 2.13 writes; its 'c' payload is identical to version 2’s.

nemo_rl.utils.cuda_ipc.normalize_cuda_ipc_handle(handle: Any) → Any[source]#

Rewrite a newer 'c'-type storage handle so older torch can open it.

Parameters:

handle – The storage_handle element of a rebuild_cuda_tensor argument tuple. Anything that is not a versioned 'c' handle from a newer torch is returned unchanged.

Returns:

The handle, with its version byte lowered to

data:

LEGACY_SHAREABLE_HANDLE_VERSION when it is a 'c' handle carrying :data:CURRENT_SHAREABLE_HANDLE_VERSION.

Raises:

ValueError –

For a 'c' handle with a version byte newer than

data:

CURRENT_SHAREABLE_HANDLE_VERSION, whose payload has not been verified against the legacy layout.