cupyx.distributed.init_process_group#

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

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

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

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

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

排名为0的进程将使用子进程启动一个TCP服务器，该服务器监听由环境变量 CUPYX_DISTRIBUTED_PORT 指定的端口，排名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