CuPy 分布式计算完全指南cupyx.distributed 进程通信与跨设备分布式数组【免费下载链接】cupyNumPy SciPy for GPU项目地址: https://gitcode.com/GitHub_Trending/cu/cupy本文基于 CuPy 仓库 docs/source/reference/distributed.rst 展开系统讲解cupyx.distributed模块的两大能力一是通过init_process_group/NCCLBackend在进程间建立基于 NCCL 的集合通信二是通过DistributedArray/distributed_array/make_2d_index_map/matmul将单个逻辑数组切分成多个 chunk 分布在多块 GPU 上并自动完成模式转换、resharding 与分布式矩阵乘法。读完本文你将掌握 CuPy 多进程通信的初始化流程、全部集合通信原语以及分布式数组的切片、模式mode与分块blocking机制能够编写可运行的多 GPU 分布式计算代码。一、cupyx.distributed 模块总览cupyx.distributed是 CuPy 中负责多进程 / 多设备计算的子包其公开接口收敛为两个层面恰好对应 docs/source/reference/distributed.rst 的两个小节层级命名空间公开 API进程间通信cupyx.distributedinit_process_group、NCCLBackend跨设备分布式数组cupyx.distributed.arraydistributed_array、DistributedArray、make_2d_index_map、matmul从入口文件可以确认这套 API 的导出关系cupyx/distributed/init.py 导出init_process_group与NCCLBackendcupyx/distributed/array/init.py 除文档列出的四个符号外还导出了Mode与REPLICA、MIN、MAX、SUM、PROD五种模式常量模式机制在第三节详述。从整体设计看这两个层面是独立可用的你可以只用init_process_group拿到 NCCL 通信子做最原始的集合通信也可以只使用DistributedArray在已初始化的多设备环境中进行分布式数组运算。二、进程间通信init_process_group 与 NCCLBackend2.1 初始化函数与参数init_process_group是分布式环境的唯一入口定义在 cupyx/distributed/_init.pyinit_process_group( n_devices, rank, *, backendnccl, hostNone, portNone, use_mpiFalse)调用它之后会返回一个符合Backend规范的通信对象。每个参与通信的进程都必须调用一次该函数。其参数含义如下参数类型必填说明n_devicesint是参与分布式执行的设备总数须满足n_devices 0rankint是当前进程关联 GPU 的唯一编号须满足0 rank n_devicesbackendstr否通信后端当前仅支持nccl默认ncclhoststr否初始化时进程会合rendezvous的地址默认None见下方环境变量portint否初始化时进程会合的端口默认None见下方环境变量use_mpibool否若为False不使用 MPI改用内置 TCP 服务器交换 CPU 侧信息默认False源码中对应有四道前置校验cupyx/distributed/_init.pyn_devices 0抛ValueError(Invalid number of devices ...)rank越界抛ValueError(Invalid number of rank ...)backend不在_backends目前只有nccl抛ValueError后端为nccl但nccl.available为假时抛RuntimeError(NCCL is not available)。2.2 会合机制TCP 服务器与两个环境变量cupyx.distributed不使用外部调度器而是手工指定 rank 并跨主机启动进程。初始化阶段会合rendezvous逻辑如下rank 0 进程启动一个子进程运行 TCP 服务器_store.TCPStore生成 NCCL unique id 写入存储然后等待其他 rank 通过 barrier 汇合非 rank 0 进程通过TCPStoreProxy连接 rank 0 的服务器读取nccl_id后再汇合。上述行为体现在 cupyx/distributed/_nccl_comm.py 的_init_with_tcp_store中。与之相关的两个环境变量定义在 cupyx/distributed/_store.py环境变量作用默认值CUPYX_DISTRIBUTED_HOSTrank 0 所在主机地址127.0.0.1CUPYX_DISTRIBUTED_PORTrank 0 监听端口13333当host/port参数为None时会依次回退到环境变量、再回退到默认值见 cupyx/distributed/_init.py。TCPStore 本身是一个轻量的 KLV 协议存储服务器仅承担初始化阶段的少量信息交换其实现不追求性能见 cupyx/distributed/_store.py 注释。如果安装了mpi4py且use_mpiTrue则改用_init_with_mpi只把 MPI 当作管理通道由 rank 0 生成nccl_id后通过MPI.COMM_WORLD.bcast分发给所有人cupyx/distributed/_nccl_comm.py。注意此时 MPI 的 rank 可以与n_devices指定的 rank 不同。2.3 使用前提与约束源码 docstring 明确了几条使用前提写代码前必须确认一个返回的 communicator 只允许绑定一个设备用户有责任在创建和使用 communicator 之前通过cupy.cuda.Device(i).use()设置好当前 GPU当前需要用户手动指定每个进程的 rank 和总进程数并在不同主机上手工启动所有进程该特性预期运行在可信集群环境中docstring 原话expected to be used within a trusted cluster environment。官方 docstring 给出了双进程的完整示例cupyx/distributed/_init.pyimport 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))两个进程分别在 GPU 0 / GPU 1 上执行process_0广播全 1 数组process_1接收后与原数组逐元素比较最终两个进程的array均为 1。2.4 NCCLBackend 的集合通信原语NCCLBackendcupyx/distributed/_nccl_comm.py是对 NCCL 的封装共提供 12 个通信方法统一签名模式为(数组..., opsum, streamNone)方法语义关键约束all_reduce(in, out, op, stream)全进程归约op∈ sum/prod/min/maxreduce(in, out, root, op, stream)归约到root进程仅 root 的out被修改broadcast(in_out, root, stream)从 root 广播非 root 的数组作为输出reduce_scatter(in, out, count, op, stream)归约后分发count为发给每个 rank 的元素数all_gather(in, out, count, stream)全 gathercount为每个 rank 的元素数send(array, peer, stream)/recv(out, peer, stream)点对点收发需成对调用send_recv(in, out, peer, stream)同时收发内部使用nccl.groupStart/groupEndscatter(in, out, root, stream)分发in_array.shape[0] n_devicesgather(in, out, root, stream)收集out_array.shape[0] n_devicesall_to_all(in, out, stream)全交换输入输出首维均为n_devicesbarrier()屏障CPU 侧显式同步会阻塞线程推进归约操作符op支持sum、prod、min、max内部映射为NCCL_SUM / NCCL_PROD / NCCL_MAX / NCCL_MIN其中复数数组只支持sum其余操作会抛ValueError(Only nccl.SUM is supported for complex arrays)见 cupyx/distributed/_nccl_comm.py。数据类型映射_nccl_dtypes定义了 NumPy dtype 字符到 NCCL 数据类型的映射cupyx/distributed/_nccl_comm.py覆盖int8/uint8/int32/uint32/int64/uint64/float16/float32/float64复数complex64(F) / complex128(D)会按实数拆分并以元素数加倍的方式传输cupyx/distributed/_nccl_comm.py。未映射的 dtype 抛TypeError。连续性约束传入的数组必须为 C 连续或 F 连续否则抛RuntimeError_check_contiguouscupyx/distributed/_nccl_comm.py。流stream支持所有方法都接受可选的stream参数传None时使用cupy.cuda.stream.get_current_stream()_get_stream。send_recv、scatter、gather、all_to_all会在内部使用nccl.groupStart() / nccl.groupEnd()将多条 NCCL 调用组合为一个 group避免死锁并提升吞吐。2.5 稠密与稀疏数组的分派NCCLBackend通过_dispatch_arg_type自动分派到稠密_DenseNCCLCommunicator或稀疏_SparseNCCLCommunicator实现cupyx/distributed/_nccl_comm.py当第一个参数是cupyx.scipy.sparse的稀疏矩阵或由稀疏矩阵组成的 list/tuple时走稀疏路径。稠密路径直接把array.data.ptr传给 NCCL 底层 API稀疏路径支持coo / csr / csc三种格式_get_sparse_type发送时把矩阵拆成内部数组coo的data/row/colcsr/csc的data/indptr/indices逐一传输归约只支持sum/prodcupyx/distributed/_nccl_comm.py。稀疏矩阵的形状和内部数组大小等元数据必须先在各进程间交换以确定接收缓冲大小——若未安装mpi4py元数据交换也会走 NCCL源码会发出警告提示这会引发设备同步和严重性能下降建议安装 MPI 与mpi4py规避cupyx/distributed/_nccl_comm.py。_Backend抽象基类cupyx/distributed/_comm.py还定义了stop()方法只有 rank 0 会关闭 TCP store。三、跨设备分布式数组DistributedArray 与 distributed_array3.1 核心概念chunk、index_map 与 modeDistributedArray继承自cupy.ndarraycupyx/distributed/array/_array.py代表一个分布在多块 CUDA 设备上的多维数组。它的核心设计有三个概念chunk块一个逻辑数组被切成若干 chunk每个 chunk 是原数组某个切片对应的连续数组。一个设备可以持有多个 chunk。index_map索引映射dict[int, list[tuple[slice, ...]]]描述每个设备 ID 拥有哪些切片的索引。mode模式决定 chunk 之间重叠区域如何解释。DistributedArray.mode目前支持五种定义在 cupyx/distributed/array/_modes.py模式常量语义幂等重叠区单位元REPLICA重叠区各 chunk 值保证一致——MIN重叠区代表原数据的min是dtype 最大值MAX重叠区代表原数据的max是dtype 最小值SUM重叠区代表原数据的sum否0PROD重叠区代表原数据的prod否1从源码看cupyx/distributed/array/_modes.pyMIN/MAX是幂等操作SUM/PROD非幂等因此非幂等模式在传输数据前会把源 chunk 的重叠区域写入单位元apply_to中的逻辑见 cupyx/distributed/array/_chunk.py。在REPLICA模式下重叠区域的值在所有 chunk 上保证一致在其他模式下重叠区域存储的是原数据经归约后的结果。大多数操作如 ufunc 和 matmul都会先把数组自动转换到合适模式因此日常使用中用户通常无需手动管理模式。3.2 工厂函数 distributed_arrayDistributedArray的直接构造函数设计为内部调用用户应使用distributed_array工厂函数创建cupyx/distributed/array/_array.pydistributed_array(array, index_map, modeREPLICA)arrayDistributedArray、cupy.ndarray或任何可传给numpy.array的对象index_mapdict[int, 切片索引]指定各设备拥有的 chunk一个设备可有多个 chunk用列表给出mode默认REPLICA。docstring 中的示例cupyx/distributed/array/_array.pyarray cupy.arange(9).reshape(3, 3) A distributed_array( array, {0: [(slice(2), slice(2)), # array[:2, :2] slice(None, None, 2)], # array[::2] 1: (slice(1, None), 2)}) # array[1:, 2]注意distributed_array不检查给定数组的所有元素是否都落在某个 chunk 中当源数组是 GPU 上的cupy.ndarray时数据会通过内部创建的 communicator 异步拷贝到各目标设备源在主机内存时则同步拷贝cupyx/distributed/array/_array.py。若传入的array本身是DistributedArray则该函数会按需执行change_mode(mode)与reshard(index_map)再返回新对象——即distributed_array也承担模式转换 重新分片的入口角色。3.3 常用属性与方法DistributedArray提供以下经过验证的属性和方法mode当前模式见上表devices持有数据的设备 ID 集合index_mapdict[int, list[切片]]各设备拥有的 chunk 索引可作为reshard的输入all_chunks()返回各设备上所有已冲刷flush缓冲数据的 chunk 数组dict[int, list[ndarray]]。resharding 和模式转换会生成缓冲数据all_chunks会先把缓冲刷入 chunk 再返回cupyx/distributed/array/_array.pychange_mode(mode)返回指定模式下的视图或副本reshard(index_map)返回具有新index_map的视图或副本。数据跨设备传输在内部创建的独立 stream上进行传输结果先缓冲、需要时再反映到 chunk 中以实现异步cupyx/distributed/array/_array.pyget()返回主机内存上的numpy.ndarray副本不支持stream/order/out参数。非REPLICA模式下未覆盖区域用对应模式单位元填充cupyx/distributed/array/_array.py。3.4 支持与暂不支持的操作支持的操作元素级运算ufunc通过__cupy_override_elementwise_kernel__钩子分发给_elementwise._execute。当多个参数共享同一index_map时走_execute_kernel逐 chunk 直接执行 kernel当index_map不同需要跨设备访问时若设备支持 P2P 则通过 peer access 直接读取否则退化为异步拷贝_execute_peer_access见 cupyx/distributed/array/_elementwise.py。沿轴的归约sum/prod/max/min通过__cupy_override_reduction_kernel__分发到_reduction._execute要求axis非None且不支持out与keepdimscupyx/distributed/array/_array.py。矩阵乘法运算符与numpy.matmul语义见第四节。shape属性可读赋值不支持。暂不支持抛出NotImplementedError__getitem__/__setitem__/__len__/__iter__、copy、reshape、transpose、astype、mean、std、var、all、any、argmax、argsort、nonzero、searchsorted、sort、ravel、repeat、clip、fill、view、dot、trace、diagonal、T/mT属性等完整清单见 cupyx/distributed/array/_array.py。这意味着当前的DistributedArray是面向特定计算场景元素级运算、归约、矩阵乘的最小实现尚不支持通用索引与就地修改。四、二维分块工具与分布式矩阵乘法4.1 make_2d_index_map按分块生成索引映射make_2d_index_map用于为二维矩阵按指定分块blocking生成index_mapcupyx/distributed/array/_linalg.pymake_2d_index_map(i_partitions, j_partitions, devices)i_partitionsi轴上的块边界列表含 0 与末尾边界j_partitionsj轴上的块边界列表deviceslen(i_partitions)-1 × len(j_partitions)-1的二维列表每个元素是拥有该块的设备 ID 集合一个块可被多设备持有。docstring 示例cupyx/distributed/array/_linalg.pyindex_map make_2d_index_map( [0, 2, 4], [0, 3, 5], [[{0}, {1}], [{2}, {0, 1}]]) # 结果 # {0: [(slice(0, 2, None), slice(0, 3, None)), # (slice(2, 4, None), slice(3, 5, None))], # 1: [(slice(0, 2, None), slice(3, 5, None)), # (slice(2, 4, None), slice(3, 5, None))], # 2: [(slice(2, 4, None), slice(0, 3, None))]}注意块(2:4, 3:5)同时被设备 0 和 1 持有——这就是一个块可被多设备持有的直接体现也是后续 matmul 执行计划可以就近选择计算设备的依据。函数要求i_partitions[0] 0、j_partitions[0] 0且分区严格递增assert校验。4.2 matmul分块执行计划与自动模式转换matmul(a, b, outNone, **kwargs)执行分布式矩阵乘法cupyx/distributed/array/_linalg.py。核心流程如下模式转换将a、b都转为REPLICA模式结果以SUM模式输出docstring 原话converts its operands into the replica mode, and compute their product in the sum mode。这是因为矩阵乘需要对共享维度做累加重叠区域天然需要 sum 语义。维度提升一维向量自动按 NumPy 规则前后补 1 维参与运算结束后再压缩掉_prepend_one_to_shape/_append_one_to_shape/_pop_from_shape。分块blocking推导由两个数组的index_map在 i / j / k 三个方向上的所有边界推导出统一的分块方案_Blocking_find_blocking要求步长为 1 且索引映射一致否则抛RuntimeError(Inconsistent index mapping)。执行计划execution plan对每个输出块(i, j)在 k 方向上寻找同时持有块a[i,k]与b[k,j]的设备作为计算设备交集非空若没有任何设备同时拥有两块抛RuntimeError(There is no device that can perform multiplication ...)cupyx/distributed/array/_linalg.py。逐块计算在每个计算设备上调用cupy.linalg._product.matmul完成局部块乘结果作为SUM模式下的新 chunk 记录并基于输入 chunk 的ready事件做流同步全程异步。docstring 中的完整示例cupyx/distributed/array/_linalg.pyA distributed_array( cupy.arange(6).reshape(2, 3), make_2d_index_map([0, 2], [0, 1, 3], [[{0}, {1, 2}]])) B distributed_array( cupy.arange(12).reshape(3, 4), make_2d_index_map([0, 1, 3], [0, 2, 4], [[{0}, {0}], [{1}, {2}]])) C A B C.mode # sum C.all_chunks() # {0: [array([[0, 0], [0, 3]]), array([[0, 0], [6, 9]])], # 1: [array([[20, 23], [56, 65]])], # 2: [array([[26, 29], [74, 83]])]} C # array([[20, 23, 26, 29], [56, 68, 80, 92]])随后可调用C.change_mode(replica)把重叠区归约回一致值模式变为replica各设备上得到完整一致的逻辑矩阵cupyx/distributed/array/_array.py。当前限制out参数、subok/axes/axis关键字均不支持抛RuntimeError混合分布式与非分布式数组也不支持抛RuntimeError(Mixing a distributed array with a non-distributed array is not supported)见 cupyx/distributed/array/_linalg.py。__matmul__与__array_ufunc__钩子会把运算符路由到该函数因此A B与matmul(A, B)等价。五、运行方式与注意事项总结5.1 最小运行步骤安装带 CUDA NCCL 的 CuPy 发行版nccl.available为真在每个参与进程内先cupy.cuda.Device(rank).use()绑定 GPU调用cupyx.distributed.init_process_group(n_devices, rank)获取通信子可选用distributed_arraymake_2d_index_map构造分布式数组进行 ufunc / 归约 /运算通过comm.stop()rank 0与进程退出释放资源。5.2 关键注意事项每个进程一个设备communicator 与设备一一对应绑定错误的 GPU 会导致通信数据错位先启动 rank 0TCP 会合机制依赖 rank 0 先行监听其他 rank 的TCPStoreProxy内置 50 次、每次 0.5 秒的重试cupyx/distributed/_store.py但长时间未就绪最终会抛RuntimeError(TCPStore is not available)可信集群前提模块未内置认证与加密不应暴露在不可信网络复数归约仅支持 sum、数组需 C/F 连续、稀疏矩阵仅 coo/csr/csc 且归约仅 sum/prodMPI 是可选加速项仅用于元数据交换与 barrier未安装时功能仍可用但稀疏通信会退化为 NCCL 传元数据并产生性能警告DistributedArray 能力边界索引[]、就地修改、mean/std/var等暂不支持设计上聚焦元素级运算、沿轴归约与矩阵乘三类场景。从源码结构看cupyx.distributed目前仍处于基础原语齐备、高级算子逐步扩充的阶段——_elementwise.py、_reduction.py、_linalg.py中的多处TODO如公平分配计算负载避免重叠数据的重复传输也印证了这一点。若要在生产环境使用建议先在单机多卡上以127.0.0.1默认会合方式验证上述最小流程再扩展到多主机集群。【免费下载链接】cupyNumPy SciPy for GPU项目地址: https://gitcode.com/GitHub_Trending/cu/cupy创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
