cupyx.distributed.NCCLBackend#

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

使用 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)[source]#

执行 all gather（全收集）操作。

参数:

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

all_reduce(in_array, out_array, op='sum', stream=None)[source]#

执行 all reduce（全规约）操作。

参数:

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)[source]#

执行 all to all（全交换）操作。

参数:

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

barrier()[source]#

执行 barrier（栅栏）操作。

栅栏操作在 CPU 中完成，是一种显式的同步机制，会暂停线程的执行。

broadcast(in_out_array, root=0, stream=None)[source]#

执行 broadcast（广播）操作。

参数:

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

gather(in_array, out_array, root=0, stream=None)[source]#

执行 gather（收集）操作。

参数:

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)[source]#

执行 receive（接收）操作。

参数:

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

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

执行 reduce（规约）操作。

参数:

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)[source]#

执行 reduce scatter（规约分散）操作。

参数:

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

scatter(in_array, out_array, root=0, stream=None)[source]#

执行 scatter（分散）操作。

参数:

in_array (cupy.ndarray) – 要发送的数组。其形状必须为 (total_ranks, …)。
out_array (cupy.ndarray) – 存储结果的数组。
root (int) – 将 in_array 发送给其他秩的秩。
stream (cupy.cuda.Stream, optional) – 如果支持，用于执行通信的流。

send(array, peer, stream=None)[source]#

执行 send（发送）操作。

参数:

array (cupy.ndarray) – 要发送的数组。
peer (int) – 将接收 array 的进程的秩。
stream (cupy.cuda.Stream, optional) – 如果支持，用于执行通信的流。

send_recv(in_array, out_array, peer, stream=None)[source]#

执行 send and receive（发送和接收）操作。

参数:

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