cupyx.distributed.init_process_group#

cupyx.distributed.init_process_group(n_devices, rank, *, backend='nccl', host=None, port=None, use_mpi=False)[source]#

启动 cupyx.distributed 并获取一个通信器。

此调用初始化分布式环境，每个参与通信的进程都需要调用它。

每个返回的通信器只允许关联一个设备。用户有责任在创建和使用通信器之前设置适当的 GPU。

当前，用户需要手动指定每个进程的 rank 和进程总数，并在不同的主机上手动启动所有进程。

rank 为 0 的进程将使用一个子进程生成一个 TCP 服务器，该服务器侦听环境变量 CUPYX_DISTRIBUTED_PORT 指定的端口，rank 0 必须在环境变量 CUPYX_DISTRIBUTED_HOST 指定的主机上执行。如果未指定这些值，则默认使用 ‘127.0.0.1’ 和 13333。

请注意，此功能应在受信任的集群环境中使用。

示例

>>> import cupy
>>> def process_0():
...     import cupyx.distributed
...     cupy.cuda.Device(0).use()
...     comm = cupyx.distributed.init_process_group(2, 0)
...     array = cupy.ones(1)
...     comm.broadcast(array, 0)
...
>>> def process_1():
...     import cupyx.distributed
...     cupy.cuda.Device(1).use()
...     comm = cupyx.distributed.init_process_group(2, 1)
...     array = cupy.zeros(1)
...     comm.broadcast(array, 0)
...     cupy.equal(array, cupy.ones(1))
参数:

  • n_devices (int) – 分布式执行中将使用的设备总数。

  • rank (int) – 通信器关联的 GPU 的唯一 ID，其值必须满足 0 <= rank < n_devices。

  • backend (str) – 用于通信的后端。可选，默认为 “nccl”。

  • host (str) – 初始化时用于进程集合的主机地址，默认为 None。

  • port (int) – 初始化时用于进程集合的端口，默认为 None。

  • use_mpi (bool) – 如果为 False，则避免使用 MPI 进行同步，并使用提供的 TCP 服务器交换仅限 CPU 的信息。默认为 False。

返回:

用于执行通信的对象，遵循

Backend 规范

返回类型:

Backend