cupyx.distributed.NCCLBackend#

class cupyx.distributed.NCCLBackend(n_devices, rank, host='127.0.0.1', port=13333, use_mpi=False)[源代码][源代码]#

使用 NVIDIA 的 NCCL 执行通信的接口。

参数:

n_devices (int) – 在分布式执行中将使用的设备总数。
rank (int) – 通信器所关联的GPU的唯一ID，其值需要满足 0 <= rank < n_devices。
host (str, optional) – 进程在初始化时集合的主机地址。默认为 “127.0.0.1”。
port (int, optional) – 用于进程在初始化时集合的端口。默认为 13333。
use_mpi (bool, optional) – 在 MPI 和使用包含的 TCP 服务器进行初始化与同步之间切换。默认为 False。

方法

all_gather(in_array, out_array, count, stream=None)[源代码][源代码]#

执行一个全收集操作。

参数:

in_array (cupy.ndarray) – 要发送的数组。
out_array (cupy.ndarray) – 存储结果的数组。
count (int) – 每个等级发送的元素数量。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

all_reduce(in_array, out_array, op='sum', stream=None)[源代码][源代码]#

执行一个全规约操作。

参数:

in_array (cupy.ndarray) – 要发送的数组。
out_array (cupy.ndarray) – 存储结果的数组。
op (str) – 归约操作，可以是 (‘sum’, ‘prod’, ‘min’ ‘max’) 之一，复数类型的数组仅支持 ‘sum’。默认为 ‘sum’。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

all_to_all(in_array, out_array, stream=None)[源代码][源代码]#

执行全对全操作。

参数:

in_array (cupy.ndarray) – 要发送的数组。其形状必须是 (total_ranks, …)。
out_array (cupy.ndarray) – 存储结果的数组。其形状必须为 (total_ranks, …)。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

barrier()[源代码][源代码]#

执行一个屏障操作。

屏障是在 CPU 中完成的，并且是一种显式的同步机制，它会暂停线程的进展。

broadcast(in_out_array, root=0, stream=None)[源代码][源代码]#

执行广播操作。

参数:

in_out_array (cupy.ndarray) – 要发送给 root 等级的数组。其他等级将在此接收广播数据。
root (int, optional) – 将发送广播的进程的等级。默认为 0。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

gather(in_array, out_array, root=0, stream=None)[源代码][源代码]#

执行一个收集操作。

参数:

in_array (cupy.ndarray) – 要发送的数组。
out_array (cupy.ndarray) – 存储结果的数组。其形状必须为 (total_ranks, …)。
root (int) – 将从其他等级接收 in_array 的等级。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

recv(out_array, peer, stream=None)[源代码][源代码]#

执行接收操作。

参数:

array (cupy.ndarray) – 用于接收数据的数组。
peer (int) – 将从接收进程 array 的等级。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

reduce(in_array, out_array, root=0, op='sum', stream=None)[源代码][源代码]#

执行一个归约操作。

参数:

in_array (cupy.ndarray) – 要发送的数组。
out_array (cupy.ndarray) – 存储结果的数组。只会被 root 进程修改。
root (int, optional) – 执行归约操作的进程的等级。默认为 0。
op (str) – 归约操作，可以是 (‘sum’, ‘prod’, ‘min’ ‘max’) 之一，复数类型的数组仅支持 ‘sum’。默认为 ‘sum’。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

reduce_scatter(in_array, out_array, count, op='sum', stream=None)[源代码][源代码]#

执行一个reduce scatter操作。

参数:

in_array (cupy.ndarray) – 要发送的数组。
out_array (cupy.ndarray) – 存储结果的数组。
count (int) – 每个等级发送的元素数量。
op (str) – 归约操作，可以是 (‘sum’, ‘prod’, ‘min’ ‘max’) 之一，复数类型的数组仅支持 ‘sum’。默认为 ‘sum’。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

scatter(in_array, out_array, root=0, stream=None)[源代码][源代码]#

执行分散操作。

参数:

in_array (cupy.ndarray) – 要发送的数组。其形状必须是 (total_ranks, …)。
out_array (cupy.ndarray) – 存储结果的数组。
root (int) – 将 in_array 发送到其他等级的等级。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

send(array, peer, stream=None)[源代码][源代码]#

执行发送操作。

参数:

array (cupy.ndarray) – 要发送的数组。
peer (int) – 进程 array 的等级将被发送至。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。

send_recv(in_array, out_array, peer, stream=None)[源代码][源代码]#

执行发送和接收操作。

参数:

in_array (cupy.ndarray) – 要发送的数组。
out_array (cupy.ndarray) – 用于接收数据的数组。
peer (int) – 发送 in_array 和接收 out_array 的进程等级。
stream (cupy.cuda.Stream, optional) – 如果支持，则使用流进行通信。