xnn.common.distributed.init_process_group

xnn.common.distributed.init_process_group(device)

Join the launcher’s process group, if there is one.

In a distributed run each rank is pinned to one CUDA device selected by LOCAL_RANK (the configured device only chooses cpu vs cuda), and the process group is initialized with the matching backend: NCCL on GPUs, Gloo on CPUs. A process group that already exists (for example created by a caller that runs several xnn components in one process) is reused.

Parameters:

device (torch.device) – The device resolved from the configuration.

Returns:

The device this rank should compute on (device unchanged in a single-process run, cuda:LOCAL_RANK or device on CPU in a distributed one), and whether this call created the process group, in which case the caller owns it and should pass the flag to destroy_process_group() when done.

Return type:

tuple of (torch.device, bool)