`nemo_rl.distributed.stateless_process_group`#

Module Contents#

Classes#

`_VllmNcclUniqueId`
`StatelessProcessGroup`

Functions#

_pickle_vllm_unique_id

Serialize an NCCL unique ID in vLLM’s metadata wire format.

Data#

`_NEMO_UNIQUE_ID_KEY`
`_VLLM_UNIQUE_ID_KEY`
`_VLLM_NCCL_MODULE`
`_VLLM_PICKLE_LOCK`

API#

nemo_rl.distributed.stateless_process_group._NEMO_UNIQUE_ID_KEY#: ‘nccl_unique_id’

nemo_rl.distributed.stateless_process_group._VLLM_UNIQUE_ID_KEY#: ‘broadcast_from/0/0’

nemo_rl.distributed.stateless_process_group._VLLM_NCCL_MODULE#: ‘vllm.distributed.device_communicators.pynccl_wrapper’

nemo_rl.distributed.stateless_process_group._VLLM_PICKLE_LOCK#: ‘Lock(…)’

class nemo_rl.distributed.stateless_process_group._VllmNcclUniqueId#

Bases: ctypes.Structure

_fields_#: [(‘internal’,)]

nemo_rl.distributed.stateless_process_group._pickle_vllm_unique_id(unique_id_bytes: bytes) → bytes#

Serialize an NCCL unique ID in vLLM’s metadata wire format.

vLLM’s stateless process group pickles its ncclUniqueId ctypes structure. Training workers do not install vLLM, so construct the same ctypes type under its canonical module name only while serializing.

class nemo_rl.distributed.stateless_process_group.StatelessProcessGroup( master_address: str, port: int, rank: int, world_size: int, )#

Initialization

abort() → None#

Terminate in-flight operations and release the communicator.

Idempotent, and safe on a group whose communicator was never built.

abort(), not destroy(), is the correct teardown here. NCCL documents destroy as an intra-node collective that every rank must call or it hangs – precisely what a rank whose process has died cannot do. abort terminates outstanding operations instead, so it works whether or not the peers are alive, which makes it the only safe choice on a path that exists to handle dead peers.

Verified on 2xA6000: with a peer SIGKILLed mid-broadcast, a survivor blocked in the collective was released 0.15s after another thread called abort().

The rendezvous store is dropped too. Each rebuild gets a fresh port, so holding the old one costs nothing functionally, but a run that recovers repeatedly would otherwise accumulate a bound TCPStore per recovery for the life of the worker.

The split children are aborted first, and they are a third communicator family. The Python reshard path splits this communicator per replica group and caches the children; NCCL gives a split child its own abort flag unless splitShare is set (it defaults to 0), so aborting this communicator does not reach them. A rank blocked on a child would never be released, and since it never returns, the watchdog’s guarded block never exits to have fired read – a hang no exception-translation fix can reach.

Imported locally to keep this module free of a weight_sync dependency at module scope.

init_nccl_communicator(device: int, *, peer: str = 'nemo') → None#

Initialize NCCL using the metadata and warmup protocol of the peer.

peer="nemo" publishes the raw 128-byte unique ID under nccl_unique_id and warms up with a rank-zero broadcast. peer="vllm" additionally publishes vLLM’s pickled ncclUniqueId under broadcast_from/0/0 and warms up with an all-reduce, matching PyNcclCommunicator. The receiver protocol is not negotiable, so a generation backend must select the peer it implements.

broadcast( tensor: torch.Tensor, src: int, stream: Optional[torch.cuda.Stream] = None, )#

nemo_rl.distributed.stateless_process_group#